CHAD let

2 files changed, 40 insertions, 25 deletions

diff --git a/src/AST.hs b/src/AST.hs
index 4d642ba..3179d92 100644
--- a/src/AST.hs
+++ b/src/AST.hs
@@ -91,6 +91,11 @@ data Expr x env t where
   EVar :: x t -> STy t -> Idx env t -> Expr x env t
   ELet :: x t -> Expr x env a -> Expr x (a : env) t -> Expr x env t
 
+  -- base types
+  EPair :: x (TPair a b) -> Expr x env a -> Expr x env b -> Expr x env (TPair a b)
+  EFst :: x a -> STy b -> Expr x env (TPair a b) -> Expr x env a
+  ESnd :: x b -> STy a -> Expr x env (TPair a b) -> Expr x env b
+
   -- array operations
   EBuild1 :: x (TArr (S Z) t) -> Expr x env TIx -> Expr x (TIx : env) t -> Expr x env (TArr (S Z) t)
   EBuild :: x (TArr n t) -> SNat n -> Vec n (Expr x env TIx) -> Expr x (ConsN n TIx env) t -> Expr x env (TArr n t)
@@ -103,7 +108,10 @@ data Expr x env t where
   EOp :: x t -> SOp a t -> Expr x env a -> Expr x env t
 
   -- EVM operations
-  EMOne :: Idx venv t -> Expr x env t -> Expr x env (TEVM venv TNil)
+  EMOne :: SList STy venv -> Idx venv t -> Expr x env t -> Expr x env (TEVM venv TNil)
+  EMScope :: Expr x env (TEVM (t : venv) a) -> Expr x env (TEVM venv (TPair a t))
+  EMReturn :: SList STy venv -> Expr x env t -> Expr x env (TEVM venv t)
+  EMBind :: Expr x env (TEVM venv a) -> Expr x (a : env) (TEVM venv b) -> Expr x env (TEVM venv b)
 deriving instance (forall ty. Show (x ty)) => Show (Expr x env t)
 
 type SOp :: Ty -> Ty -> Type
@@ -141,7 +149,7 @@ typeOf = \case
   EIdx _ e _ | STArr _ t <- typeOf e -> t
   EOp _ op _ -> opt2 op
 
-  EMOne _ _ -> STEVM _ STNil
+  EMOne t _ _ -> STEVM t STNil
 
 unSNat :: SNat n -> Nat
 unSNat SZ = Z
@@ -175,6 +183,7 @@ data env :> env' where
   WId :: env :> env
   WSink :: env :> (t : env)
   WCopy :: env :> env' -> (t : env) :> (t : env')
+  WPop :: (t : env) :> env' -> env :> env'
   WThen :: env1 :> env2 -> env2 :> env3 -> env1 :> env3
 deriving instance Show (env :> env')
 
@@ -187,6 +196,7 @@ WId @> i = i
 WSink @> i = IS i
 WCopy _ @> IZ = IZ
 WCopy w @> (IS i) = IS (w @> i)
+WPop w @> i = w @> IS i
 WThen w1 w2 @> i = w2 @> w1 @> i
 
 weakenExpr :: env :> env' -> Expr x env t -> Expr x env' t
@@ -200,7 +210,7 @@ weakenExpr w = \case
   EIdx1 x e1 e2 -> EIdx1 x (weakenExpr w e1) (weakenExpr w e2)
   EIdx x e1 es -> EIdx x (weakenExpr w e1) (weakenVec w es)
   EOp x op e -> EOp x op (weakenExpr w e)
-  EMOne i e -> EMOne i (weakenExpr w e)
+  EMOne t i e -> EMOne t i (weakenExpr w e)
 
 wcopyN :: SNat n -> env :> env' -> ConsN n TIx env :> ConsN n TIx env'
 wcopyN SZ w = w
diff --git a/src/CHAD.hs b/src/CHAD.hs
index 17ee12b..f01ab9e 100644
--- a/src/CHAD.hs
+++ b/src/CHAD.hs
@@ -37,14 +37,9 @@ bconcat b1 BTop = b1
 bconcat b1 (BPush b2 x) = BPush (bconcat b1 b2) x
 
 bconcat' :: (forall e1 e2 t. e1 :> e2 -> f e1 t -> f e2 t)
-         -> Bindings f env env1 -> Bindings f env env2 -> (forall env12. Bindings env env12 -> r) -> r
-bconcat' wf b1 b2 = weakenBindings
--- bconcat :: (forall e1 e2 t. e1 :> e2 -> f e1 t -> f e2 t)
---         -> Bindings f env env1 -> Bindings f env env2 -> env :> env'
---         -> (forall env'12. Bindings f env' env'12 -> r) -> r
--- bconcat wf BTop b w k = weakenBindings wf w b $ \b' _ -> k b'
--- bconcat wf (BPush b x) b2 w k =
---   bconcat wf 
+         -> Bindings f env env1 -> Bindings f env env2
+         -> (forall env12. Bindings f env env12 -> r) -> r
+bconcat' wf b1 b2 k = weakenBindings wf (sinkWithBindings b1) b2 $ \b2' _ -> k (bconcat b1 b2')
 
 type family D1 t where
   D1 TNil = TNil
@@ -73,13 +68,6 @@ type family D2E env where
   D2E '[] = '[]
   D2E (t : env) = D2 t : D2E env
 
-data Ret env t =
-  forall env'.
-    Ret (Bindings Ex (D1E env) env')
-        (Ex env' (D1 t))
-        (Ex (D2 t : env') (TEVM (D2E env) TNil))
-deriving instance Show (Ret env t)
-
 d1 :: STy t -> STy (D1 t)
 d1 STNil = STNil
 d1 (STPair a b) = STPair (d1 a) (d1 b)
@@ -99,6 +87,10 @@ d2 (STScal t) = case t of
   STBool -> STNil
 d2 STEVM{} = error "EVM not allowed in input program"
 
+d2list :: SList STy env -> SList STy (D2E env)
+d2list SNil = SNil
+d2list (SCons x l) = SCons (d2 x) (d2list l)
+
 conv1Idx :: Idx env t -> Idx (D1E env) (D1 t)
 conv1Idx IZ = IZ
 conv1Idx (IS i) = IS (conv1Idx i)
@@ -107,15 +99,28 @@ conv2Idx :: Idx env t -> Idx (D2E env) (D2 t)
 conv2Idx IZ = IZ
 conv2Idx (IS i) = IS (conv2Idx i)
 
-drev :: Ex env t -> Ret env t
-drev = \case
+data Ret env t =
+  forall env'.
+    Ret (Bindings Ex (D1E env) env')
+        (Ex env' (D1 t))
+        (Ex (D2 t : env') (TEVM (D2E env) TNil))
+deriving instance Show (Ret env t)
+
+drev :: SList STy env -> Ex env t -> Ret env t
+drev senv = \case
   EVar _ t i ->
     Ret BTop
         (EVar ext (d1 t) (conv1Idx i))
-        (EMOne (conv2Idx i) (EVar ext (d2 t) IZ))
-  ELet _ rhs body ->
-    Ret _
-        _
-        _
+        (EMOne (d2list senv) (conv2Idx i) (EVar ext (d2 t) IZ))
+  ELet _ rhs body
+    | Ret rhs0 rhs1 rhs2 <- drev senv rhs
+    , Ret body0 body1 body2 <- drev (SCons (typeOf rhs) senv) body ->
+    weakenBindings weakenExpr (WCopy (sinkWithBindings rhs0)) body0 $ \body0' wbody0' ->
+    Ret (bconcat (BPush rhs0 rhs1) body0')
+        (weakenExpr wbody0' body1)
+        (EMBind (EMScope (weakenExpr (WCopy wbody0') body2))
+                (ELet ext (ESnd ext STNil (EVar ext (STPair STNil (d2 (typeOf rhs))) IZ)) $
+                  weakenExpr (WCopy (WSink .> WSink .> WPop (sinkWithBindings body0'))) rhs2))
+  _ -> undefined
   where
     ext = Const ()