
   0 | module Language.Reflection.VarSubst
   1 | 
   2 | import Data.SortedSet
   3 | import Data.SortedMap
   4 | import Data.Maybe
   5 | import Language.Reflection
   6 | import Language.Reflection.TTImp
   7 | import Language.Reflection.TT
   8 | import Language.Reflection.Util
   9 | import Control.Monad.Identity
  10 | import Control.Monad.Reader
  11 | import Control.Monad.Reader.Tuple
  12 | import Control.Monad.State
  13 | 
  14 | %default total
  15 | 
  16 | ||| Monadic operation on a TTImp
  17 | TTOp : (Type -> Type) -> Type
  18 | TTOp m = TTImp -> m TTImp -> m TTImp
  19 | 
  20 | ||| SortedSet of Names
  21 | public export
  22 | NameSet : Type
  23 | NameSet = SortedSet Name
  24 | 
  25 | ||| Insert `Maybe Name`'s content into a NameSet
  26 | insertM : Maybe Name -> NameSet -> NameSet
  27 | insertM Nothing = id
  28 | insertM (Just x) = insert x
  29 | 
  30 | ||| Information about variable shadowing
  31 | record ShadowingInfo where
  32 |   constructor MkSI
  33 |   ||| Set of shadowed names
  34 |   shadowedNames : NameSet
  35 | 
  36 | ||| Check if a variable is shadowed
  37 | isShadowed : Name -> ShadowingInfo -> Bool
  38 | isShadowed n = contains n . shadowedNames
  39 | 
  40 | parameters {auto _ : MonadReader ShadowingInfo m}
  41 |            (f : TTOp m)
  42 |   mutual
  43 |     ||| Provide a ShadowingInfo for an operation on a TTImp over a PiInfo TTImp
  44 |     doPiInfo : PiInfo TTImp -> m $ PiInfo TTImp
  45 |     doPiInfo (DefImplicit x) = DefImplicit <$> doTTImp x
  46 |     doPiInfo x = pure x
  47 | 
  48 |     ||| Provide a ShadowingInfo for an operation on a TTImp
  49 |     doTTImp : TTImp -> m TTImp
  50 |     doTTImp = assert_total mapATTImp' provideSI
  51 | 
  52 |     ||| Provide a ShadowingInfo for an operation on a TTImp
  53 |     provideSI : TTOp m
  54 |     provideSI b@(IPi fc count pinfo mn argTy lamTy) newM = do
  55 |       pinfo <- doPiInfo pinfo
  56 |       argTy <- doTTImp argTy
  57 |       lamTy <- local {shadowedNames $= insertM mn} $ doTTImp lamTy
  58 |       f b $ pure $ IPi fc count pinfo mn argTy lamTy
  59 |     provideSI b@(ILam fc count pinfo mn argTy lamTy) newM = do
  60 |       pinfo <- doPiInfo pinfo
  61 |       argTy <- doTTImp argTy
  62 |       lamTy <- local {shadowedNames $= insertM mn} $ doTTImp lamTy
  63 |       f b $ pure $ ILam fc count pinfo mn argTy lamTy
  64 |     provideSI b@(ILet fc lhsFC c n nTy nVal scope) newM = do
  65 |       nTy <- doTTImp nTy
  66 |       nVal <- doTTImp nVal
  67 |       scope <- local {shadowedNames $= insert n} $ doTTImp scope
  68 |       f b $ pure $ ILet fc lhsFC c n nTy nVal scope
  69 |     provideSI b newM = do
  70 |       f b newM
  71 | 
  72 | ||| Information about quoting
  73 | record QuoteInfo where
  74 |   constructor MkQI
  75 |   ||| Indicator of being in a quote
  76 |   isQuote : Bool
  77 | 
  78 | ||| Provide QuoteInfo for a monadic operation on a TTImp
  79 | provideQI : MonadReader QuoteInfo m =>
  80 |             (f : TTOp m) ->
  81 |             TTOp m
  82 | provideQI f b@(IQuote fc t) newM = f b $ local {isQuote := True} newM
  83 | provideQI f b@(IQuoteDecl fc decls) newM = f b $ local {isQuote := True} newM
  84 | provideQI f b@(IUnquote fc t) newM = f b $ local {isQuote := False} newM
  85 | provideQI f b newM = f b newM
  86 | 
  87 | ||| Provide QuoteInfo and ShadowingInfo for a monadic operation on a TTImp
  88 | provideCtx : Monad m =>
  89 |              MonadReader QuoteInfo m =>
  90 |              MonadReader ShadowingInfo m =>
  91 |              TTOp m -> TTOp m
  92 | provideCtx = provideSI . provideQI
  93 | 
  94 | ||| Run a monadic operation requiring ShadowingInfo and QuoteInfo on a TTImp
  95 | mapTTOp : Monad m => TTOp (ReaderT (ShadowingInfo, QuoteInfo) m) -> TTImp -> m TTImp
  96 | mapTTOp op t = do
  97 |   let readerOp = mapATTImp' . provideCtx $ op
  98 |   runReaderT (MkSI empty, MkQI False) $ readerOp t
  99 | 
 100 | containsVariableImpl : Monad m =>
 101 |                        MonadReader QuoteInfo m =>
 102 |                        MonadReader ShadowingInfo m =>
 103 |                        MonadState Bool m =>
 104 |                        Name -> TTOp m
 105 | containsVariableImpl n (IVar _ n') m =
 106 |   if n == n'
 107 |      && not !(asks isQuote)
 108 |      && not !(asks $ isShadowed n)
 109 |   then put True *> m
 110 |   else m
 111 | containsVariableImpl n tt m = m
 112 | 
 113 | ||| Check if a TTImp contains a variable
 114 | export
 115 | containsVariable : Name -> TTImp -> Bool
 116 | containsVariable n =
 117 |   execState False .
 118 |     mapTTOp (containsVariableImpl n)
 119 | 
 120 | collectVariablesImpl : Monad m =>
 121 |                        MonadReader QuoteInfo m =>
 122 |                        MonadReader ShadowingInfo m =>
 123 |                        MonadState NameSet m =>
 124 |                        TTOp m
 125 | collectVariablesImpl (IVar _ n) m = do
 126 |   if not !(asks isQuote)
 127 |      && not !(asks $ isShadowed n)
 128 |      then modify (insert n) *> m
 129 |      else m
 130 | collectVariablesImpl tt m = m
 131 | 
 132 | ||| Calculate a set of used variables
 133 | export
 134 | usesVariables : TTImp -> NameSet
 135 | usesVariables =
 136 |   execState empty .
 137 |     mapTTOp (collectVariablesImpl)
 138 | 
 139 | substituteVariablesImpl : Monad m =>
 140 |                           MonadReader QuoteInfo m =>
 141 |                           MonadReader ShadowingInfo m =>
 142 |                           SortedMap Name TTImp -> TTOp m
 143 | substituteVariablesImpl vMap (IVar fc n) m =
 144 |   if not !(asks isQuote)
 145 |      && not !(asks $ isShadowed n)
 146 |   then
 147 |     fromMaybe m $ pure <$> lookup n vMap
 148 |   else m
 149 | substituteVariablesImpl _ _ m = m
 150 | 
 151 | ||| Substitute all occurrences of each variable with the given expression
 152 | export
 153 | substituteVariables : SortedMap Name TTImp -> TTImp -> TTImp
 154 | substituteVariables vMap =
 155 |   runIdentity .
 156 |     mapTTOp (substituteVariablesImpl vMap)
 157 | 
 158 | ||| Substitute all occurrences of given variable with given expression
 159 | export
 160 | substituteVariable : Name -> TTImp -> TTImp -> TTImp
 161 | substituteVariable n t =
 162 |   runIdentity .
 163 |     mapTTOp (substituteVariablesImpl $ insert n t empty)
 164 | 
 165 | 
 166 | substituteBindImpl : Monad m =>
 167 |                      MonadReader QuoteInfo m =>
 168 |                      MonadReader ShadowingInfo m =>
 169 |                      SortedMap Name TTImp -> TTOp m
 170 | substituteBindImpl vMap (IVar fc n) m =
 171 |   if not !(asks isQuote)
 172 |      && not !(asks $ isShadowed n)
 173 |   then
 174 |     fromMaybe m $ pure <$> lookup n vMap
 175 |   else m
 176 | substituteBindImpl vMap (IBindVar fc n) m =
 177 |   if not !(asks isQuote)
 178 |      && not !(asks $ isShadowed n)
 179 |   then
 180 |     fromMaybe m $ pure <$> lookup n vMap
 181 |   else m
 182 | substituteBindImpl _ _ m = m
 183 | 
 184 | ||| Substitute all variable and BindVar occurrences with given
 185 | export
 186 | substituteBind : SortedMap Name TTImp -> TTImp -> TTImp
 187 | substituteBind vMap = do
 188 |   runIdentity . mapTTOp (substituteBindImpl vMap)
 189 |