1 files changed, 47 insertions, 0 deletions
diff --git a/Sink.hs b/Sink.hs
new file mode 100644
index 0000000..1368cb6
--- /dev/null
+++ b/Sink.hs
@@ -0,0 +1,47 @@
+{-# LANGUAGE DataKinds #-}
+{-# LANGUAGE GADTs #-}
+{-# LANGUAGE LambdaCase #-}
+{-# LANGUAGE RankNTypes #-}
+{-# LANGUAGE TypeOperators #-}
+module Sink where
+
+import AST
+
+
+newtype env :> env' = Weaken { (>:>) :: forall t'. Idx env t' -> Idx env' t' }
+
+wId :: env :> env
+wId = Weaken id
+
+wSucc :: env :> env' -> env :> (a ': env')
+wSucc (Weaken f) = Weaken (Succ . f)
+
+wSink :: env :> env' -> (a ': env) :> (a ': env')
+wSink w = Weaken (\case Zero -> Zero
+                        Succ i -> Succ (w >:> i))
+
+(.>) :: env2 :> env3 -> env1 :> env2 -> env1 :> env3
+Weaken f .> Weaken g = Weaken (f . g)
+
+sinkExp :: env :> env' -> Exp env a -> Exp env' a
+sinkExp w = \case
+    App e1 e2 -> App (sinkExp w e1) (sinkExp w e2)
+    Lam t e -> Lam t (sinkExp (wSink w) e)
+    Var t i -> Var t (w >:> i)
+    Let e1 e2 -> Let (sinkExp w e1) (sinkExp (wSink w) e2)
+    Lit l -> Lit l
+    Cond e1 e2 e3 -> Cond (sinkExp w e1) (sinkExp w e2) (sinkExp w e3)
+    Const c -> Const c
+    Pair e1 e2 -> Pair (sinkExp w e1) (sinkExp w e2)
+    Fst e -> Fst (sinkExp w e)
+    Snd e -> Snd (sinkExp w e)
+    Build sht e1 e2 -> Build sht (sinkExp w e1) (sinkExp w e2)
+    Ifold sht e1 e2 e3 -> Ifold sht (sinkExp w e1) (sinkExp w e2) (sinkExp w e3)
+    Index e1 e2 -> Index (sinkExp w e1) (sinkExp w e2)
+    Shape e -> Shape (sinkExp w e)
+
+sinkExp1 :: Exp env a -> Exp (t ': env) a
+sinkExp1 = sinkExp (wSucc wId)
+
+sinkExp2 :: Exp env a -> Exp (t1 ': t2 ': env) a
+sinkExp2 = sinkExp (wSucc (wSucc wId))