module Cryptol.IR.FreeVars
  ( FreeVars(..)
  , Deps(..)
  , Defs(..)
  , moduleDeps, transDeps
  ) where

import           Data.Set ( Set )
import qualified Data.Set as Set
import           Data.Map ( Map )
import qualified Data.Map as Map

import Cryptol.TypeCheck.AST
import Cryptol.Utils.RecordMap

data Deps = Deps { Deps -> Set Name
valDeps  :: Set Name
                   -- ^ Undefined value names

                 , Deps -> Set Name
tyDeps   :: Set Name
                   -- ^ Undefined type names (from newtype)

                 , Deps -> Set TParam
tyParams :: Set TParam
                   -- ^ Undefined type params (e.d. mod params)
                 } deriving Deps -> Deps -> Bool
(Deps -> Deps -> Bool) -> (Deps -> Deps -> Bool) -> Eq Deps
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Deps -> Deps -> Bool
$c/= :: Deps -> Deps -> Bool
== :: Deps -> Deps -> Bool
$c== :: Deps -> Deps -> Bool
Eq

instance Semigroup Deps where
  Deps
d1 <> :: Deps -> Deps -> Deps
<> Deps
d2 = [Deps] -> Deps
forall a. Monoid a => [a] -> a
mconcat [Deps
d1,Deps
d2]

instance Monoid Deps where
  mempty :: Deps
mempty = Deps :: Set Name -> Set Name -> Set TParam -> Deps
Deps { valDeps :: Set Name
valDeps   = Set Name
forall a. Set a
Set.empty
                , tyDeps :: Set Name
tyDeps    = Set Name
forall a. Set a
Set.empty
                , tyParams :: Set TParam
tyParams  = Set TParam
forall a. Set a
Set.empty
                }

  mappend :: Deps -> Deps -> Deps
mappend Deps
d1 Deps
d2 = Deps
d1 Deps -> Deps -> Deps
forall a. Semigroup a => a -> a -> a
<> Deps
d2

  mconcat :: [Deps] -> Deps
mconcat [Deps]
ds = Deps :: Set Name -> Set Name -> Set TParam -> Deps
Deps { valDeps :: Set Name
valDeps   = [Set Name] -> Set Name
forall (f :: * -> *) a. (Foldable f, Ord a) => f (Set a) -> Set a
Set.unions ((Deps -> Set Name) -> [Deps] -> [Set Name]
forall a b. (a -> b) -> [a] -> [b]
map Deps -> Set Name
valDeps [Deps]
ds)
                    , tyDeps :: Set Name
tyDeps    = [Set Name] -> Set Name
forall (f :: * -> *) a. (Foldable f, Ord a) => f (Set a) -> Set a
Set.unions ((Deps -> Set Name) -> [Deps] -> [Set Name]
forall a b. (a -> b) -> [a] -> [b]
map Deps -> Set Name
tyDeps  [Deps]
ds)
                    , tyParams :: Set TParam
tyParams  = [Set TParam] -> Set TParam
forall (f :: * -> *) a. (Foldable f, Ord a) => f (Set a) -> Set a
Set.unions ((Deps -> Set TParam) -> [Deps] -> [Set TParam]
forall a b. (a -> b) -> [a] -> [b]
map Deps -> Set TParam
tyParams [Deps]
ds)
                    }

rmTParam :: TParam -> Deps -> Deps
rmTParam :: TParam -> Deps -> Deps
rmTParam TParam
p Deps
x = Deps
x { tyParams :: Set TParam
tyParams = TParam -> Set TParam -> Set TParam
forall a. Ord a => a -> Set a -> Set a
Set.delete TParam
p (Deps -> Set TParam
tyParams Deps
x) }

rmVal :: Name -> Deps -> Deps
rmVal :: Name -> Deps -> Deps
rmVal Name
p Deps
x = Deps
x { valDeps :: Set Name
valDeps = Name -> Set Name -> Set Name
forall a. Ord a => a -> Set a -> Set a
Set.delete Name
p (Deps -> Set Name
valDeps Deps
x) }

rmVals :: Set Name -> Deps -> Deps
rmVals :: Set Name -> Deps -> Deps
rmVals Set Name
p Deps
x = Deps
x { valDeps :: Set Name
valDeps = Set Name -> Set Name -> Set Name
forall a. Ord a => Set a -> Set a -> Set a
Set.difference (Deps -> Set Name
valDeps Deps
x) Set Name
p }


-- | Compute the transitive closure of the given dependencies.
transDeps :: Map Name Deps -> Map Name Deps
transDeps :: Map Name Deps -> Map Name Deps
transDeps Map Name Deps
mp0 = (Map Name Deps, Map Name Deps) -> Map Name Deps
forall a b. (a, b) -> a
fst
              ((Map Name Deps, Map Name Deps) -> Map Name Deps)
-> (Map Name Deps, Map Name Deps) -> Map Name Deps
forall a b. (a -> b) -> a -> b
$ [(Map Name Deps, Map Name Deps)] -> (Map Name Deps, Map Name Deps)
forall a. [a] -> a
head
              ([(Map Name Deps, Map Name Deps)]
 -> (Map Name Deps, Map Name Deps))
-> [(Map Name Deps, Map Name Deps)]
-> (Map Name Deps, Map Name Deps)
forall a b. (a -> b) -> a -> b
$ ((Map Name Deps, Map Name Deps) -> Bool)
-> [(Map Name Deps, Map Name Deps)]
-> [(Map Name Deps, Map Name Deps)]
forall a. (a -> Bool) -> [a] -> [a]
dropWhile ((Map Name Deps -> Map Name Deps -> Bool)
-> (Map Name Deps, Map Name Deps) -> Bool
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry Map Name Deps -> Map Name Deps -> Bool
forall a. Eq a => a -> a -> Bool
(/=))
              ([(Map Name Deps, Map Name Deps)]
 -> [(Map Name Deps, Map Name Deps)])
-> [(Map Name Deps, Map Name Deps)]
-> [(Map Name Deps, Map Name Deps)]
forall a b. (a -> b) -> a -> b
$ [Map Name Deps]
-> [Map Name Deps] -> [(Map Name Deps, Map Name Deps)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Map Name Deps]
steps ([Map Name Deps] -> [Map Name Deps]
forall a. [a] -> [a]
tail [Map Name Deps]
steps)
  where
  step1 :: Map Name Deps -> Deps -> Deps
step1 Map Name Deps
mp Deps
d = [Deps] -> Deps
forall a. Monoid a => [a] -> a
mconcat [ Deps -> Name -> Map Name Deps -> Deps
forall k a. Ord k => a -> k -> Map k a -> a
Map.findWithDefault
                            Deps
forall a. Monoid a => a
mempty { valDeps :: Set Name
valDeps = Name -> Set Name
forall a. a -> Set a
Set.singleton Name
x }
                            Name
x Map Name Deps
mp | Name
x <- Set Name -> [Name]
forall a. Set a -> [a]
Set.toList (Deps -> Set Name
valDeps Deps
d) ]
  step :: Map Name Deps -> Map Name Deps
step Map Name Deps
mp = (Deps -> Deps) -> Map Name Deps -> Map Name Deps
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Map Name Deps -> Deps -> Deps
step1 Map Name Deps
mp) Map Name Deps
mp

  steps :: [Map Name Deps]
steps = (Map Name Deps -> Map Name Deps)
-> Map Name Deps -> [Map Name Deps]
forall a. (a -> a) -> a -> [a]
iterate Map Name Deps -> Map Name Deps
step Map Name Deps
mp0

-- | Dependencies of top-level declarations in a module.
-- These are dependencies on module parameters or things
-- defined outside the module.
moduleDeps :: Module -> Map Name Deps
moduleDeps :: Module -> Map Name Deps
moduleDeps = Map Name Deps -> Map Name Deps
transDeps (Map Name Deps -> Map Name Deps)
-> (Module -> Map Name Deps) -> Module -> Map Name Deps
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Map Name Deps] -> Map Name Deps
forall (f :: * -> *) k a.
(Foldable f, Ord k) =>
f (Map k a) -> Map k a
Map.unions ([Map Name Deps] -> Map Name Deps)
-> (Module -> [Map Name Deps]) -> Module -> Map Name Deps
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (DeclGroup -> Map Name Deps) -> [DeclGroup] -> [Map Name Deps]
forall a b. (a -> b) -> [a] -> [b]
map DeclGroup -> Map Name Deps
forall e. (Defs e, FreeVars e) => e -> Map Name Deps
fromDG ([DeclGroup] -> [Map Name Deps])
-> (Module -> [DeclGroup]) -> Module -> [Map Name Deps]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Module -> [DeclGroup]
mDecls
  where
  fromDG :: e -> Map Name Deps
fromDG e
dg = let vs :: Deps
vs = e -> Deps
forall e. FreeVars e => e -> Deps
freeVars e
dg
              in [(Name, Deps)] -> Map Name Deps
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList [ (Name
x,Deps
vs) | Name
x <- Set Name -> [Name]
forall a. Set a -> [a]
Set.toList (e -> Set Name
forall d. Defs d => d -> Set Name
defs e
dg) ]

class FreeVars e where
  freeVars :: e -> Deps


instance FreeVars e => FreeVars [e] where
  freeVars :: [e] -> Deps
freeVars = [Deps] -> Deps
forall a. Monoid a => [a] -> a
mconcat ([Deps] -> Deps) -> ([e] -> [Deps]) -> [e] -> Deps
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (e -> Deps) -> [e] -> [Deps]
forall a b. (a -> b) -> [a] -> [b]
map e -> Deps
forall e. FreeVars e => e -> Deps
freeVars


instance FreeVars DeclGroup where
  freeVars :: DeclGroup -> Deps
freeVars DeclGroup
dg = case DeclGroup
dg of
                  NonRecursive Decl
d -> Decl -> Deps
forall e. FreeVars e => e -> Deps
freeVars Decl
d
                  Recursive [Decl]
ds   -> Set Name -> Deps -> Deps
rmVals ([Decl] -> Set Name
forall d. Defs d => d -> Set Name
defs [Decl]
ds) ([Decl] -> Deps
forall e. FreeVars e => e -> Deps
freeVars [Decl]
ds)


instance FreeVars Decl where
  freeVars :: Decl -> Deps
freeVars Decl
d = DeclDef -> Deps
forall e. FreeVars e => e -> Deps
freeVars (Decl -> DeclDef
dDefinition Decl
d) Deps -> Deps -> Deps
forall a. Semigroup a => a -> a -> a
<> Schema -> Deps
forall e. FreeVars e => e -> Deps
freeVars (Decl -> Schema
dSignature Decl
d)


instance FreeVars DeclDef where
  freeVars :: DeclDef -> Deps
freeVars DeclDef
d = case DeclDef
d of
                 DeclDef
DPrim -> Deps
forall a. Monoid a => a
mempty
                 DExpr Expr
e -> Expr -> Deps
forall e. FreeVars e => e -> Deps
freeVars Expr
e


instance FreeVars Expr where
  freeVars :: Expr -> Deps
freeVars Expr
expr =
    case Expr
expr of
      EList [Expr]
es Type
t        -> [Expr] -> Deps
forall e. FreeVars e => e -> Deps
freeVars [Expr]
es Deps -> Deps -> Deps
forall a. Semigroup a => a -> a -> a
<> Type -> Deps
forall e. FreeVars e => e -> Deps
freeVars Type
t
      ETuple [Expr]
es         -> [Expr] -> Deps
forall e. FreeVars e => e -> Deps
freeVars [Expr]
es
      ERec RecordMap Ident Expr
fs           -> [Expr] -> Deps
forall e. FreeVars e => e -> Deps
freeVars (RecordMap Ident Expr -> [Expr]
forall a b. RecordMap a b -> [b]
recordElements RecordMap Ident Expr
fs)
      ESel Expr
e Selector
_          -> Expr -> Deps
forall e. FreeVars e => e -> Deps
freeVars Expr
e
      ESet Type
ty Expr
e Selector
_ Expr
v     -> Type -> Deps
forall e. FreeVars e => e -> Deps
freeVars Type
ty Deps -> Deps -> Deps
forall a. Semigroup a => a -> a -> a
<> [Expr] -> Deps
forall e. FreeVars e => e -> Deps
freeVars [Expr
e,Expr
v]
      EIf Expr
e1 Expr
e2 Expr
e3      -> [Expr] -> Deps
forall e. FreeVars e => e -> Deps
freeVars [Expr
e1,Expr
e2,Expr
e3]
      EComp Type
t1 Type
t2 Expr
e [[Match]]
mss -> [Type] -> Deps
forall e. FreeVars e => e -> Deps
freeVars [Type
t1,Type
t2] Deps -> Deps -> Deps
forall a. Semigroup a => a -> a -> a
<> Set Name -> Deps -> Deps
rmVals ([[Match]] -> Set Name
forall d. Defs d => d -> Set Name
defs [[Match]]
mss) (Expr -> Deps
forall e. FreeVars e => e -> Deps
freeVars Expr
e)
                                            Deps -> Deps -> Deps
forall a. Semigroup a => a -> a -> a
<> [Deps] -> Deps
forall a. Monoid a => [a] -> a
mconcat (([Match] -> Deps) -> [[Match]] -> [Deps]
forall a b. (a -> b) -> [a] -> [b]
map [Match] -> Deps
forall a. (FreeVars a, Defs a) => [a] -> Deps
foldFree [[Match]]
mss)
      EVar Name
x            -> Deps
forall a. Monoid a => a
mempty { valDeps :: Set Name
valDeps = Name -> Set Name
forall a. a -> Set a
Set.singleton Name
x }
      ETAbs TParam
a Expr
e         -> TParam -> Deps -> Deps
rmTParam TParam
a (Expr -> Deps
forall e. FreeVars e => e -> Deps
freeVars Expr
e)
      ETApp Expr
e Type
t         -> Expr -> Deps
forall e. FreeVars e => e -> Deps
freeVars Expr
e Deps -> Deps -> Deps
forall a. Semigroup a => a -> a -> a
<> Type -> Deps
forall e. FreeVars e => e -> Deps
freeVars Type
t
      EApp Expr
e1 Expr
e2        -> [Expr] -> Deps
forall e. FreeVars e => e -> Deps
freeVars [Expr
e1,Expr
e2]
      EAbs Name
x Type
t Expr
e        -> Type -> Deps
forall e. FreeVars e => e -> Deps
freeVars Type
t Deps -> Deps -> Deps
forall a. Semigroup a => a -> a -> a
<> Name -> Deps -> Deps
rmVal Name
x (Expr -> Deps
forall e. FreeVars e => e -> Deps
freeVars Expr
e)
      EProofAbs Type
p Expr
e     -> Type -> Deps
forall e. FreeVars e => e -> Deps
freeVars Type
p Deps -> Deps -> Deps
forall a. Semigroup a => a -> a -> a
<> Expr -> Deps
forall e. FreeVars e => e -> Deps
freeVars Expr
e
      EProofApp Expr
e       -> Expr -> Deps
forall e. FreeVars e => e -> Deps
freeVars Expr
e
      EWhere Expr
e [DeclGroup]
ds       -> [DeclGroup] -> Deps
forall a. (FreeVars a, Defs a) => [a] -> Deps
foldFree [DeclGroup]
ds Deps -> Deps -> Deps
forall a. Semigroup a => a -> a -> a
<> Set Name -> Deps -> Deps
rmVals ([DeclGroup] -> Set Name
forall d. Defs d => d -> Set Name
defs [DeclGroup]
ds) (Expr -> Deps
forall e. FreeVars e => e -> Deps
freeVars Expr
e)
    where
      foldFree :: (FreeVars a, Defs a) => [a] -> Deps
      foldFree :: [a] -> Deps
foldFree = (a -> Deps -> Deps) -> Deps -> [a] -> Deps
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr a -> Deps -> Deps
forall d. (FreeVars d, Defs d) => d -> Deps -> Deps
updateFree Deps
forall a. Monoid a => a
mempty
      updateFree :: d -> Deps -> Deps
updateFree d
x Deps
rest = d -> Deps
forall e. FreeVars e => e -> Deps
freeVars d
x Deps -> Deps -> Deps
forall a. Semigroup a => a -> a -> a
<> Set Name -> Deps -> Deps
rmVals (d -> Set Name
forall d. Defs d => d -> Set Name
defs d
x) Deps
rest

instance FreeVars Match where
  freeVars :: Match -> Deps
freeVars Match
m = case Match
m of
                 From Name
_ Type
t1 Type
t2 Expr
e -> Type -> Deps
forall e. FreeVars e => e -> Deps
freeVars Type
t1 Deps -> Deps -> Deps
forall a. Semigroup a => a -> a -> a
<> Type -> Deps
forall e. FreeVars e => e -> Deps
freeVars Type
t2 Deps -> Deps -> Deps
forall a. Semigroup a => a -> a -> a
<> Expr -> Deps
forall e. FreeVars e => e -> Deps
freeVars Expr
e
                 Let Decl
d          -> Decl -> Deps
forall e. FreeVars e => e -> Deps
freeVars Decl
d



instance FreeVars Schema where
  freeVars :: Schema -> Deps
freeVars Schema
s = (TParam -> Deps -> Deps) -> Deps -> [TParam] -> Deps
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr TParam -> Deps -> Deps
rmTParam ([Type] -> Deps
forall e. FreeVars e => e -> Deps
freeVars (Schema -> [Type]
sProps Schema
s) Deps -> Deps -> Deps
forall a. Semigroup a => a -> a -> a
<> Type -> Deps
forall e. FreeVars e => e -> Deps
freeVars (Schema -> Type
sType Schema
s))
                              (Schema -> [TParam]
sVars Schema
s)

instance FreeVars Type where
  freeVars :: Type -> Deps
freeVars Type
ty =
    case Type
ty of
      TCon TCon
tc [Type]
ts -> TCon -> Deps
forall e. FreeVars e => e -> Deps
freeVars TCon
tc Deps -> Deps -> Deps
forall a. Semigroup a => a -> a -> a
<> [Type] -> Deps
forall e. FreeVars e => e -> Deps
freeVars [Type]
ts
      TVar TVar
tv -> TVar -> Deps
forall e. FreeVars e => e -> Deps
freeVars TVar
tv

      TUser Name
_ [Type]
_ Type
t -> Type -> Deps
forall e. FreeVars e => e -> Deps
freeVars Type
t
      TRec RecordMap Ident Type
fs     -> [Type] -> Deps
forall e. FreeVars e => e -> Deps
freeVars (RecordMap Ident Type -> [Type]
forall a b. RecordMap a b -> [b]
recordElements RecordMap Ident Type
fs)


instance FreeVars TVar where
  freeVars :: TVar -> Deps
freeVars TVar
tv = case TVar
tv of
                  TVBound TParam
p -> Deps
forall a. Monoid a => a
mempty { tyParams :: Set TParam
tyParams = TParam -> Set TParam
forall a. a -> Set a
Set.singleton TParam
p }
                  TVar
_         -> Deps
forall a. Monoid a => a
mempty

instance FreeVars TCon where
  freeVars :: TCon -> Deps
freeVars TCon
tc =
    case TCon
tc of
     TC (TCNewtype (UserTC Name
n Kind
_)) -> Deps
forall a. Monoid a => a
mempty { tyDeps :: Set Name
tyDeps = Name -> Set Name
forall a. a -> Set a
Set.singleton Name
n }
     TCon
_ -> Deps
forall a. Monoid a => a
mempty

instance FreeVars Newtype where
  freeVars :: Newtype -> Deps
freeVars Newtype
nt = (TParam -> Deps -> Deps) -> Deps -> [TParam] -> Deps
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr TParam -> Deps -> Deps
rmTParam Deps
base (Newtype -> [TParam]
ntParams Newtype
nt)
    where base :: Deps
base = [Type] -> Deps
forall e. FreeVars e => e -> Deps
freeVars (Newtype -> [Type]
ntConstraints Newtype
nt) Deps -> Deps -> Deps
forall a. Semigroup a => a -> a -> a
<> [Type] -> Deps
forall e. FreeVars e => e -> Deps
freeVars (((Ident, Type) -> Type) -> [(Ident, Type)] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map (Ident, Type) -> Type
forall a b. (a, b) -> b
snd (Newtype -> [(Ident, Type)]
ntFields Newtype
nt))


--------------------------------------------------------------------------------

class Defs d where
  defs :: d -> Set Name

instance Defs a => Defs [a] where
  defs :: [a] -> Set Name
defs = [Set Name] -> Set Name
forall (f :: * -> *) a. (Foldable f, Ord a) => f (Set a) -> Set a
Set.unions ([Set Name] -> Set Name) -> ([a] -> [Set Name]) -> [a] -> Set Name
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> Set Name) -> [a] -> [Set Name]
forall a b. (a -> b) -> [a] -> [b]
map a -> Set Name
forall d. Defs d => d -> Set Name
defs

instance Defs DeclGroup where
  defs :: DeclGroup -> Set Name
defs DeclGroup
dg = case DeclGroup
dg of
              Recursive [Decl]
ds   -> [Decl] -> Set Name
forall d. Defs d => d -> Set Name
defs [Decl]
ds
              NonRecursive Decl
d -> Decl -> Set Name
forall d. Defs d => d -> Set Name
defs Decl
d

instance Defs Decl where
  defs :: Decl -> Set Name
defs Decl
d = Name -> Set Name
forall a. a -> Set a
Set.singleton (Decl -> Name
dName Decl
d)

instance Defs Match where
  defs :: Match -> Set Name
defs Match
m = case Match
m of
             From Name
x Type
_ Type
_ Expr
_ -> Name -> Set Name
forall a. a -> Set a
Set.singleton Name
x
             Let Decl
d -> Decl -> Set Name
forall d. Defs d => d -> Set Name
defs Decl
d