FreeVars.hs

{-# LANGUAGE InstanceSigs #-}
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.TypeCheck.FFI.FFIType
import Cryptol.Utils.Panic ( panic )
import Cryptol.Utils.RecordMap

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

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

                 , tyParams :: Set TParam
                   -- ^ Undefined type params (e.d. mod params)
                 } deriving Eq

instance Semigroup Deps where
  d1 <> d2 = mconcat [d1,d2]

instance Monoid Deps where
  mempty = Deps { valDeps   = Set.empty
                , tyDeps    = Set.empty
                , tyParams  = Set.empty
                }

  mappend = (<>)

  mconcat ds = Deps { valDeps   = Set.unions (map valDeps ds)
                    , tyDeps    = Set.unions (map tyDeps  ds)
                    , tyParams  = Set.unions (map tyParams ds)
                    }

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

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

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


-- | Compute the transitive closure of the given dependencies.
transDeps :: Map Name Deps -> Map Name Deps
transDeps mp0 = fst
              $ headInfList
              $ dropWhile (uncurry (/=))
              $ zip steps
              $ drop 1 steps
  where
  step1 mp d = mconcat [ Map.findWithDefault
                            mempty { valDeps = Set.singleton x }
                            x mp | x <- Set.toList (valDeps d) ]
  step mp = fmap (step1 mp) mp

  steps = iterate step mp0

  -- The only call site for this function is on the result of an invocation
  -- of `iterate`, which returns an infinite list. As such, it is safe to call
  -- `head` on it. Alternatively, we could depend on a library such as
  -- `infinite-list`, but this is likely overkill for this one use site.
  headInfList l =
    case l of
      x:_ -> x
      [] -> panic "transDeps" ["`iterate` returned empty list"]

-- | 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 = transDeps . Map.unions . map fromDG . mDecls
  where
  fromDG dg = let vs = freeVars dg
              in Map.fromList [ (x,vs) | x <- Set.toList (defs dg) ]

class FreeVars e where
  freeVars :: e -> Deps


instance FreeVars e => FreeVars [e] where
  freeVars = mconcat . map freeVars


instance FreeVars DeclGroup where
  freeVars dg = case dg of
                  NonRecursive d -> freeVars d
                  Recursive ds   -> rmVals (defs ds) (freeVars ds)


instance FreeVars Decl where
  freeVars d = freeVars (dDefinition d) <> freeVars (dSignature d)


instance FreeVars DeclDef where
  freeVars :: DeclDef -> Deps
  freeVars d = case d of
                 DPrim -> mempty
                 DForeign t me -> freeVars t <> foldMap freeVars me
                 DExpr e -> freeVars e

instance FreeVars FFI where
  freeVars c =
    case c of
      CallC _ -> mempty
      CallAbstract t ->
        let vs = freeVars (ffiArgTypes t) <> freeVars (ffiRetType t)
        in foldr rmTParam vs (ffiTParams t)

instance FreeVars Expr where
  freeVars expr =
    case expr of
      ELocated _r t     -> freeVars t
      EList es t        -> freeVars es <> freeVars t
      ETuple es         -> freeVars es
      ERec fs           -> freeVars (recordElements fs)
      ESel e _          -> freeVars e
      ESet ty e _ v     -> freeVars ty <> freeVars [e,v]
      EIf e1 e2 e3      -> freeVars [e1,e2,e3]
      ECase e as d      -> freeVars e <> freeVars (Map.elems as)
                                      <> maybe mempty freeVars d
      EComp t1 t2 e mss -> freeVars [t1,t2] <> rmVals (defs mss) (freeVars e)
                                            <> mconcat (map foldFree mss)
      EVar x            -> mempty { valDeps = Set.singleton x }
      ETAbs a e         -> rmTParam a (freeVars e)
      ETApp e t         -> freeVars e <> freeVars t
      EApp e1 e2        -> freeVars [e1,e2]
      EAbs x t e        -> freeVars t <> rmVal x (freeVars e)
      EProofAbs p e     -> freeVars p <> freeVars e
      EProofApp e       -> freeVars e
      EWhere e ds       -> foldFree ds <> rmVals (defs ds) (freeVars e)
      EPropGuards guards t -> mconcat [ freeVars e | (_, e) <- guards ]
                              <> freeVars t
    where
      foldFree :: (FreeVars a, Defs a) => [a] -> Deps
      foldFree = foldr updateFree mempty
      updateFree x rest = freeVars x <> rmVals (defs x) rest

instance FreeVars CaseAlt where
  freeVars (CaseAlt xs e) = foldr (rmVal . fst) (freeVars e) xs

instance FreeVars Match where
  freeVars m = case m of
                 From _ t1 t2 e -> freeVars t1 <> freeVars t2 <> freeVars e
                 Let d          -> freeVars d



instance FreeVars Schema where
  freeVars s = foldr rmTParam (freeVars (sProps s) <> freeVars (sType s))
                              (sVars s)

instance FreeVars Type where
  freeVars ty =
    case ty of
      TCon tc ts -> freeVars tc <> freeVars ts
      TVar tv -> freeVars tv

      TUser _ _ t -> freeVars t
      TRec fs     -> freeVars (recordElements fs)
      TNominal nt ts -> mempty { tyDeps = Set.singleton (ntName nt) }
                        <> freeVars ts

instance FreeVars TVar where
  freeVars tv = case tv of
                  TVBound p -> mempty { tyParams = Set.singleton p }
                  _         -> mempty

instance FreeVars TCon where
  freeVars _tc = mempty

instance FreeVars NominalType where
  freeVars nt = foldr rmTParam base (ntParams nt)
    where base = freeVars (ntConstraints nt) <> freeVars (ntDef nt)

instance FreeVars NominalTypeDef where
  freeVars def =
    case def of
      Struct c -> freeVars c
      Enum cs -> freeVars cs
      Abstract -> mempty

instance FreeVars StructCon where
  freeVars c = freeVars (recordElements (ntFields c))

instance FreeVars EnumCon where
  freeVars c = freeVars (ecFields c)


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

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

instance Defs a => Defs [a] where
  defs = Set.unions . map defs

instance Defs DeclGroup where
  defs dg = case dg of
              Recursive ds   -> defs ds
              NonRecursive d -> defs d

instance Defs Decl where
  defs d = Set.singleton (dName d)

instance Defs Match where
  defs m = case m of
             From x _ _ _ -> Set.singleton x
             Let d -> defs d