0 | module Core.Transform

  2 | import Core.Context

  3 | import Core.Env

  5 | import Libraries.Data.NameMap

  7 | %default total

  9 | unload : List (FC, Term vars) -> Term vars -> Term vars

 10 | unload [] fn = fn

 11 | unload ((fc, arg) :: args) fn = unload args (App fc fn arg)

 14 | data MatchVars : Scope -> Scope -> Type where

 15 |      None : MatchVars lhsvars vs

 16 |      Match : (idx : Nat) -> (0 p : IsVar n idx lhsvars) -> Term vs ->

 17 |              MatchVars lhsvars vs -> MatchVars lhsvars vs

 19 | lookupMatch : (idx : Nat) -> (0 p : IsVar n idx lhsvars) -> MatchVars lhsvars vs ->

 20 |               Maybe (Term vs)

 21 | lookupMatch idx p None = Nothing

 22 | lookupMatch idx p (Match v _ val rest)

 23 |     = if idx == v

 24 |          then Just val

 25 |          else lookupMatch idx p rest

 27 | addMatch : (idx : Nat) -> (0 p : IsVar n idx lhsvars) -> Term vs ->

 28 |            MatchVars lhsvars vs -> Maybe (MatchVars lhsvars vs)

 29 | addMatch idx p val ms

 30 |     = case lookupMatch idx p ms of

 31 |            Nothing => Just (Match idx p val ms)

 32 |            Just val' => if eqTerm val val'

 33 |                            then Just ms

 34 |                            else Nothing

 38 | match : MatchVars vars vs ->

 39 |         Term vars -> Term vs -> Maybe (MatchVars vars vs)

 40 | match ms (Local _ _ idx p) val

 41 |     = addMatch idx p val ms

 42 | match ms (App _ f a) (App _ f' a')

 43 |     = do ms' <- match ms f f'

 44 |          match ms' a a'

 45 | match ms x y

 46 |     = if eqTerm x y

 47 |          then Just ms

 48 |          else Nothing

 51 | tryReplace : MatchVars vars vs -> Term vars -> Maybe (Term vs)

 52 | tryReplace ms (Local _ _ idx p) = lookupMatch idx p ms

 53 | tryReplace ms (Ref fc nt n) = pure (Ref fc nt n)

 54 | tryReplace ms (Meta fc n i as)

 55 |     = do as' <- traverse (tryReplace ms) as

 56 |          pure (Meta fc n i as')

 57 | tryReplace ms (Bind fc x b sc)

 58 |     = Nothing -- TODO: can't do this yet... need to be able to weaken ms

 59 |               -- Rules are unlikely to have binders usually but we should

 60 |               -- still support it eventually

 61 | tryReplace ms (App fc f a)

 62 |     = do f' <- tryReplace ms f

 63 |          a' <- tryReplace ms a

 64 |          pure (App fc f' a')

 65 | tryReplace ms (As fc s a p)

 66 |     = do a' <- tryReplace ms a

 67 |          p' <- tryReplace ms p

 68 |          pure (As fc s a' p')

 69 | tryReplace ms (TDelayed fc r tm)

 70 |     = do tm' <- tryReplace ms tm

 71 |          pure (TDelayed fc r tm')

 72 | tryReplace ms (TDelay fc r ty tm)

 73 |     = do ty' <- tryReplace ms ty

 74 |          tm' <- tryReplace ms tm

 75 |          pure (TDelay fc r ty' tm')

 76 | tryReplace ms (TForce fc r tm)

 77 |     = do tm' <- tryReplace ms tm

 78 |          pure (TForce fc r tm')

 79 | tryReplace ms (PrimVal fc c) = pure (PrimVal fc c)

 80 | tryReplace ms (Erased fc Impossible) = pure (Erased fc Impossible)

 81 | tryReplace ms (Erased fc Placeholder) = pure (Erased fc Placeholder)

 82 | tryReplace ms (Erased fc (Dotted t)) = Erased fc . Dotted <$> tryReplace ms t

 83 | tryReplace ms (TType fc u) = pure (TType fc u)

 86 | tryApply : Transform -> Term vs -> Maybe (Term vs)

 87 | tryApply trans@(MkTransform {vars} n _ lhs rhs) tm

 88 |    = case match None lhs tm of

 89 |           Just ms => tryReplace ms rhs

 90 |           Nothing =>

 91 |             case tm of

 92 |                  App fc f a =>

 93 |                      do f' <- tryApply trans f

 94 |                         Just (App fc f' a)

 95 |                  _ => Nothing

 98 | apply : List Transform -> Term vars -> (Bool, Term vars)

 99 | apply [] tm = (False, tm)

100 | apply (t :: ts) tm

101 |     = case tryApply t tm of

102 |            Nothing => apply ts tm

103 |            Just res => (True, res)

105 | data Upd : Type where

108 | trans : {auto c : Ref Ctxt Defs} ->

109 |         {auto u : Ref Upd Bool} ->

110 |         Env Term vars -> List (FC, Term vars) -> Term vars ->

111 |         Core (Term vars)

112 | trans env stk (Ref fc Func fn)

113 |     = do defs <- get Ctxt

114 |          case lookup fn (transforms defs) of

115 |               Nothing => pure (unload stk (Ref fc Func fn))

116 |               Just ts => do let fullapp = unload stk (Ref fc Func fn)

117 |                             let (u, tm') = apply ts fullapp

118 |                             update Upd (|| u)

119 |                             pure tm'

120 | trans env stk (Meta fc n i args)

121 |     = do args' <- traverse (trans env []) args

122 |          pure $ unload stk (Meta fc n i args')

123 | trans env stk (Bind fc x b sc)

124 |     = do b' <- traverse (trans env []) b

125 |          sc' <- trans (b' :: env) [] sc

126 |          pure $ unload stk (Bind fc x b' sc')

127 | trans env stk (App fc fn arg)

128 |     = do arg' <- trans env [] arg

129 |          trans env ((fc, arg') :: stk) fn

130 | trans env stk (TDelayed fc r tm)

131 |     = do tm' <- trans env [] tm

132 |          pure $ unload stk (TDelayed fc r tm')

133 | trans env stk (TDelay fc r ty tm)

134 |     = do ty' <- trans env [] ty

135 |          tm' <- trans env [] tm

136 |          pure $ unload stk (TDelay fc r ty' tm')

137 | trans env stk (TForce fc r tm)

138 |     = do tm' <- trans env [] tm

139 |          pure $ unload stk (TForce fc r tm')

140 | trans env stk tm = pure $ unload stk tm

143 | transLoop : {auto c : Ref Ctxt Defs} ->

144 |             Nat -> Env Term vars -> Term vars -> Core (Term vars)

145 | transLoop Z env tm = pure tm

146 | transLoop (S k) env tm

147 |     = do u <- newRef Upd False

148 |          tm' <- trans env [] tm

149 |          upd <- get Upd

150 |          if upd -- If there was a transform applied, go around again until

151 |                 -- we hit the threshold

152 |             then transLoop k env tm'

153 |             else pure tm'

157 | applyTransforms : {auto c : Ref Ctxt Defs} ->

158 |                   Env Term vars -> Term vars -> Core (Term vars)

159 | applyTransforms env tm = transLoop 5 env tm