{-# LANGUAGE LambdaCase #-}

-- |
-- Module      : Jikka.Converter.Core.FreeVars
-- Description : provides utilities aboud free variables. / 自由変数についてのユーティリティを提供します。
-- Copyright   : (c) Kimiyuki Onaka, 2020
-- License     : Apache License 2.0
-- Maintainer  : kimiyuki95@gmail.com
-- Stability   : experimental
-- Portability : portable
module Jikka.Core.Language.FreeVars where

import Data.Maybe
import qualified Data.Set as S
import Jikka.Core.Language.Expr

freeVars :: Expr -> S.Set VarName
freeVars :: Expr -> Set VarName
freeVars = \case
  Var VarName
x -> VarName -> Set VarName
forall a. a -> Set a
S.singleton VarName
x
  Lit Literal
_ -> Set VarName
forall a. Set a
S.empty
  App Expr
f Expr
e -> Expr -> Set VarName
freeVars Expr
f Set VarName -> Set VarName -> Set VarName
forall a. Semigroup a => a -> a -> a
<> Expr -> Set VarName
freeVars Expr
e
  Lam VarName
x Type
_ Expr
e -> VarName -> Set VarName -> Set VarName
forall a. Ord a => a -> Set a -> Set a
S.delete VarName
x (Expr -> Set VarName
freeVars Expr
e)
  Let VarName
x Type
_ Expr
e1 Expr
e2 -> Expr -> Set VarName
freeVars Expr
e1 Set VarName -> Set VarName -> Set VarName
forall a. Semigroup a => a -> a -> a
<> VarName -> Set VarName -> Set VarName
forall a. Ord a => a -> Set a -> Set a
S.delete VarName
x (Expr -> Set VarName
freeVars Expr
e2)
  Assert Expr
e1 Expr
e2 -> Expr -> Set VarName
freeVars Expr
e1 Set VarName -> Set VarName -> Set VarName
forall a. Semigroup a => a -> a -> a
<> Expr -> Set VarName
freeVars Expr
e2

-- | `isFreeVar` checks if the given variable occurs in the tiven expr. This considers contexts.
--
-- >>> VarName "x" `isFreeVar` Lam (VarName "y") IntTy (Var (VarName "x"))
-- True
--
-- >>> VarName "x" `isFreeVar` Lam (VarName "x") IntTy (Var (VarName "x"))
-- False
isFreeVar :: VarName -> Expr -> Bool
isFreeVar :: VarName -> Expr -> Bool
isFreeVar VarName
x = \case
  Var VarName
y -> VarName
y VarName -> VarName -> Bool
forall a. Eq a => a -> a -> Bool
== VarName
x
  Lit Literal
_ -> Bool
False
  App Expr
f Expr
e -> VarName -> Expr -> Bool
isFreeVar VarName
x Expr
f Bool -> Bool -> Bool
|| VarName -> Expr -> Bool
isFreeVar VarName
x Expr
e
  Lam VarName
y Type
_ Expr
e -> VarName
x VarName -> VarName -> Bool
forall a. Eq a => a -> a -> Bool
/= VarName
y Bool -> Bool -> Bool
&& VarName -> Expr -> Bool
isFreeVar VarName
x Expr
e
  Let VarName
y Type
_ Expr
e1 Expr
e2 -> (VarName
y VarName -> VarName -> Bool
forall a. Eq a => a -> a -> Bool
/= VarName
x Bool -> Bool -> Bool
&& VarName -> Expr -> Bool
isFreeVar VarName
x Expr
e1) Bool -> Bool -> Bool
|| VarName -> Expr -> Bool
isFreeVar VarName
x Expr
e2
  Assert Expr
e1 Expr
e2 -> VarName -> Expr -> Bool
isFreeVar VarName
x Expr
e1 Bool -> Bool -> Bool
|| VarName -> Expr -> Bool
isFreeVar VarName
x Expr
e2

-- | `isUnusedVar` is the negation of `isFreeVar`.
--
-- TODO: rename to `isNonFreeVar`?
isUnusedVar :: VarName -> Expr -> Bool
isUnusedVar :: VarName -> Expr -> Bool
isUnusedVar VarName
x Expr
e = Bool -> Bool
not (VarName -> Expr -> Bool
isFreeVar VarName
x Expr
e)

-- | `isFreeVarOrScopedVar` checks if the given variable occurs in the tiven expr. This ignores contexts.
--
-- >>> VarName "x" `isFreeVarOrScopedVar` Lam (VarName "x") IntTy (Var (VarName "y"))
-- True
isFreeVarOrScopedVar :: VarName -> Expr -> Bool
isFreeVarOrScopedVar :: VarName -> Expr -> Bool
isFreeVarOrScopedVar VarName
x = \case
  Var VarName
y -> VarName
y VarName -> VarName -> Bool
forall a. Eq a => a -> a -> Bool
== VarName
x
  Lit Literal
_ -> Bool
False
  App Expr
f Expr
e -> VarName -> Expr -> Bool
isFreeVarOrScopedVar VarName
x Expr
f Bool -> Bool -> Bool
|| VarName -> Expr -> Bool
isFreeVarOrScopedVar VarName
x Expr
e
  Lam VarName
y Type
_ Expr
e -> VarName
x VarName -> VarName -> Bool
forall a. Eq a => a -> a -> Bool
== VarName
y Bool -> Bool -> Bool
|| VarName -> Expr -> Bool
isFreeVarOrScopedVar VarName
x Expr
e
  Let VarName
y Type
_ Expr
e1 Expr
e2 -> VarName
y VarName -> VarName -> Bool
forall a. Eq a => a -> a -> Bool
== VarName
x Bool -> Bool -> Bool
|| VarName -> Expr -> Bool
isFreeVarOrScopedVar VarName
x Expr
e1 Bool -> Bool -> Bool
|| VarName -> Expr -> Bool
isFreeVarOrScopedVar VarName
x Expr
e2
  Assert Expr
e1 Expr
e2 -> VarName -> Expr -> Bool
isFreeVarOrScopedVar VarName
x Expr
e1 Bool -> Bool -> Bool
|| VarName -> Expr -> Bool
isFreeVarOrScopedVar VarName
x Expr
e2

freeTyVars :: Type -> [TypeName]
freeTyVars :: Type -> [TypeName]
freeTyVars = \case
  VarTy TypeName
x -> [TypeName
x]
  Type
IntTy -> []
  Type
BoolTy -> []
  ListTy Type
t -> Type -> [TypeName]
freeTyVars Type
t
  TupleTy [Type]
ts -> (Type -> [TypeName]) -> [Type] -> [TypeName]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Type -> [TypeName]
freeTyVars [Type]
ts
  FunTy Type
t1 Type
t2 -> Type -> [TypeName]
freeTyVars Type
t1 [TypeName] -> [TypeName] -> [TypeName]
forall a. [a] -> [a] -> [a]
++ Type -> [TypeName]
freeTyVars Type
t2
  DataStructureTy DataStructure
_ -> []

findUnusedVarName :: VarName -> Expr -> VarName
findUnusedVarName :: VarName -> Expr -> VarName
findUnusedVarName (VarName OccName
x NameFlavour
_) Expr
e =
  let xs :: Set Int
xs = [Int] -> Set Int
forall a. Ord a => [a] -> Set a
S.fromList ((VarName -> [Int]) -> [VarName] -> [Int]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (\(VarName OccName
_ NameFlavour
i) -> NameFlavour -> [Int]
forall a. Maybe a -> [a]
maybeToList NameFlavour
i) (Set VarName -> [VarName]
forall a. Set a -> [a]
S.toList (Expr -> Set VarName
freeVars Expr
e)))
      flavour :: Int
flavour = [Int] -> Int
forall a. [a] -> a
head ([Int] -> Int) -> [Int] -> Int
forall a b. (a -> b) -> a -> b
$ (Int -> Bool) -> [Int] -> [Int]
forall a. (a -> Bool) -> [a] -> [a]
filter (Int -> Set Int -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.notMember` Set Int
xs) [Int
0 ..]
   in OccName -> NameFlavour -> VarName
VarName OccName
x (Int -> NameFlavour
forall a. a -> Maybe a
Just Int
flavour)

findUnusedVarName' :: Expr -> VarName
findUnusedVarName' :: Expr -> VarName
findUnusedVarName' = VarName -> Expr -> VarName
findUnusedVarName (OccName -> NameFlavour -> VarName
VarName OccName
forall a. Maybe a
Nothing NameFlavour
forall a. Maybe a
Nothing)