91 lines
3.3 KiB
Agda
91 lines
3.3 KiB
Agda
module Lib.LiftWhere
|
||
|
||
import Prelude
|
||
import Lib.Common
|
||
import Lib.Types
|
||
import Lib.TopContext
|
||
import Lib.Ref2
|
||
import Data.SortedMap
|
||
import Data.IORef
|
||
|
||
-- track quantity, depth, and whether we need to replace Bnd with a top level call.
|
||
-- We _could_ keep track of the support and elide that as well.
|
||
LiftEnv : U
|
||
LiftEnv = List (Quant × Maybe (QName × List Quant))
|
||
|
||
liftWhereTm : {{Ref2 Defs St}} → QName → LiftEnv → Tm → M Tm
|
||
liftWhereTm name env (Lam fc nm icit quant t) =
|
||
Lam fc nm icit quant <$> liftWhereTm name ((quant, Nothing) :: env) t
|
||
liftWhereTm name env (App fc t u) =
|
||
App fc <$> liftWhereTm name env t <*> liftWhereTm name env u
|
||
liftWhereTm name env (Pi fc nm icit quant t u) = do
|
||
t <- liftWhereTm name env t
|
||
u <- liftWhereTm name ((quant, Nothing) :: env) u
|
||
pure $ Pi fc nm icit quant t u
|
||
liftWhereTm name env (Let fc nm v sc) = do
|
||
v <- liftWhereTm name env v
|
||
sc <- liftWhereTm name ((Many, Nothing) :: env) sc
|
||
pure $ Let fc nm v sc
|
||
liftWhereTm name env tm@(Case fc t alts) = do
|
||
t <- liftWhereTm name env t
|
||
alts' <- traverse liftWhereAlt alts
|
||
pure $ Case fc t alts'
|
||
where
|
||
-- This is where I wish I had put indexes on things
|
||
liftWhereAlt : CaseAlt → M CaseAlt
|
||
liftWhereAlt (CaseDefault tm) = CaseDefault <$> liftWhereTm name env tm
|
||
liftWhereAlt (CaseLit l tm) = CaseLit l <$> liftWhereTm name env tm
|
||
liftWhereAlt (CaseCons qn args tm) =
|
||
-- TODO - can we pull quantities off of these arguments?
|
||
-- We should switch to SnocList first, the args are backwards here
|
||
CaseCons qn args <$> liftWhereTm name ((map (const (Many, Nothing)) args) ++ env) tm
|
||
-- This is where the magic happens
|
||
liftWhereTm name env (LetRec fc nm ty t u) = do
|
||
-- let l = length env
|
||
qn <- getName name nm
|
||
-- I'll erase this one, since we're replacing all of the usages
|
||
let qs = Zero :: map fst env
|
||
let env' = ((Zero, Just (qn, qs)) :: env)
|
||
-- environment should subst this function (see next case)
|
||
t' <- liftWhereTm qn env' t
|
||
-- TODO we could have subst in this var and dropped the extra argument
|
||
modifyRef Defs (updateMap qn (Fn $ wrapLam qs t'))
|
||
|
||
-- The rest
|
||
u' <- liftWhereTm qn env' u
|
||
pure $ LetRec fc nm (Erased fc) (Erased fc) u'
|
||
where
|
||
-- TODO might be nice if we could nest the names (Foo.go.go) for nested where
|
||
getName : QName → String → M QName
|
||
getName qn@(QN ns nm) ext = do
|
||
let qn' = QN ns (nm ++ "." ++ ext)
|
||
top <- getRef Defs
|
||
let (Just _) = lookupMap qn' top | _ => pure qn'
|
||
getName qn (ext ++ "'")
|
||
|
||
wrapLam : List Quant → Tm → Tm
|
||
wrapLam Nil tm = tm
|
||
wrapLam (q :: qs) tm = wrapLam qs $ Lam fc "_" Explicit q tm
|
||
|
||
-- And where it lands
|
||
liftWhereTm name env tm@(Bnd fc k) = case getAt (cast k) env of
|
||
Just (_, Just (qn, v)) => pure $ apply (length' env) (length' v) (Ref fc qn)
|
||
_ => pure tm
|
||
where
|
||
apply : Int → Int → Tm → Tm
|
||
apply l 0 tm = tm
|
||
-- (l - k) is like lvl2ix, but with k one bigger
|
||
apply l k tm = App fc (apply l (k - 1) tm) (Bnd fc (l - k))
|
||
liftWhereTm name env tm = pure tm
|
||
|
||
liftWhereFn : {{Ref2 Defs St}} → QName × Def → M Unit
|
||
liftWhereFn (name, Fn tm) = do
|
||
tm' <- liftWhereTm name Nil tm
|
||
modifyRef Defs $ updateMap name (Fn tm')
|
||
liftWhereFn _ = pure MkUnit
|
||
|
||
liftWhere : {{Ref2 Defs St}} → M Unit
|
||
liftWhere = do
|
||
defs <- getRef Defs
|
||
ignore $ traverse liftWhereFn $ toList defs
|