1 files changed, 45 insertions, 15 deletions

diff --git a/src/AST.hs b/src/AST.hs
index 3179d92..b1f3e5d 100644
--- a/src/AST.hs
+++ b/src/AST.hs
@@ -11,11 +11,13 @@
 {-# LANGUAGE PolyKinds #-}
 {-# LANGUAGE DeriveFunctor #-}
 {-# LANGUAGE DeriveFoldable #-}
-module AST where
+module AST (module AST, module AST.Weaken) where
 
 import Data.Kind (Type)
 import Data.Int
 
+import AST.Weaken
+
 
 data Nat = Z | S Nat
   deriving (Show, Eq, Ord)
@@ -36,10 +38,12 @@ data SList f l where
   SNil :: SList f '[]
   SCons :: f a -> SList f l -> SList f (a : l)
 deriving instance (forall a. Show (f a)) => Show (SList f l)
+infixr `SCons`
 
 data Ty
   = TNil
   | TPair Ty Ty
+  | TEither Ty Ty
   | TArr Nat Ty  -- ^ rank, element type
   | TScal ScalTy
   | TEVM [Ty] Ty
@@ -52,6 +56,7 @@ type STy :: Ty -> Type
 data STy t where
   STNil :: STy TNil
   STPair :: STy a -> STy b -> STy (TPair a b)
+  STEither :: STy a -> STy b -> STy (TEither a b)
   STArr :: SNat n -> STy t -> STy (TArr n t)
   STScal :: SScalTy t -> STy (TScal t)
   STEVM :: SList STy env -> STy t -> STy (TEVM env t)
@@ -93,8 +98,12 @@ data Expr x env t where
 
   -- 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
+  EFst :: x a -> Expr x env (TPair a b) -> Expr x env a
+  ESnd :: x b -> Expr x env (TPair a b) -> Expr x env b
+  ENil :: x TNil -> Expr x env TNil
+  EInl :: x (TEither a b) -> STy b -> Expr x env a -> Expr x env (TEither a b)
+  EInr :: x (TEither a b) -> STy a -> Expr x env b -> Expr x env (TEither a b)
+  ECase :: x c -> Expr x env (TEither a b) -> Expr x (a : env) c -> Expr x (b : env) c -> Expr x env c
 
   -- array operations
   EBuild1 :: x (TArr (S Z) t) -> Expr x env TIx -> Expr x (TIx : env) t -> Expr x env (TArr (S Z) t)
@@ -112,6 +121,9 @@ data Expr x env t where
   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)
+
+  -- partiality
+  EError :: STy a -> String -> Expr x env a
 deriving instance (forall ty. Show (x ty)) => Show (Expr x env t)
 
 type SOp :: Ty -> Ty -> Type
@@ -139,17 +151,30 @@ typeOf :: Expr x env t -> STy t
 typeOf = \case
   EVar _ t _ -> t
   ELet _ _ e -> typeOf e
+
+  EPair _ a b -> STPair (typeOf a) (typeOf b)
+  EFst _ e | STPair t _ <- typeOf e -> t
+  ESnd _ e | STPair _ t <- typeOf e -> t
+  ENil _ -> STNil
+  EInl _ t2 e -> STEither (typeOf e) t2
+  EInr _ t1 e -> STEither t1 (typeOf e)
+  ECase _ _ a _ -> typeOf a
+
   EBuild1 _ _ e -> STArr (SS SZ) (typeOf e)
   EBuild _ n _ e -> STArr n (typeOf e)
   EFold1 _ _ e | STArr (SS n) t <- typeOf e -> STArr n t
 
-  -- expression operations
   EConst _ t _ -> STScal t
   EIdx1 _ e _ | STArr (SS n) t <- typeOf e -> STArr n t
   EIdx _ e _ | STArr _ t <- typeOf e -> t
   EOp _ op _ -> opt2 op
 
   EMOne t _ _ -> STEVM t STNil
+  EMScope e | STEVM (SCons t env) a <- typeOf e -> STEVM env (STPair a t)
+  EMReturn env e -> STEVM env (typeOf e)
+  EMBind _ e -> typeOf e
+
+  EError t _ -> t
 
 unSNat :: SNat n -> Nat
 unSNat SZ = Z
@@ -159,6 +184,7 @@ unSTy :: STy t -> Ty
 unSTy = \case
   STNil -> TNil
   STPair a b -> TPair (unSTy a) (unSTy b)
+  STEither a b -> TEither (unSTy a) (unSTy b)
   STArr n t -> TArr (unSNat n) (unSTy t)
   STScal t -> TScal (unSScalTy t)
   STEVM l t -> TEVM (unSList l) (unSTy t)
@@ -179,17 +205,6 @@ fromNat :: Nat -> Int
 fromNat Z = 0
 fromNat (S n) = succ (fromNat n)
 
-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')
-
-(.>) :: env2 :> env3 -> env1 :> env2 -> env1 :> env3
-(.>) = flip WThen
-
 infixr @>
 (@>) :: env :> env' -> Idx env t -> Idx env' t
 WId @> i = i
@@ -203,6 +218,13 @@ weakenExpr :: env :> env' -> Expr x env t -> Expr x env' t
 weakenExpr w = \case
   EVar x t i -> EVar x t (w @> i)
   ELet x rhs body -> ELet x (weakenExpr w rhs) (weakenExpr (WCopy w) body)
+  EPair x e1 e2 -> EPair x (weakenExpr w e1) (weakenExpr w e2)
+  EFst x e -> EFst x (weakenExpr w e)
+  ESnd x e -> ESnd x (weakenExpr w e)
+  ENil x -> ENil x
+  EInl x t e -> EInl x t (weakenExpr w e)
+  EInr x t e -> EInr x t (weakenExpr w e)
+  ECase x e1 e2 e3 -> ECase x (weakenExpr w e1) (weakenExpr (WCopy w) e2) (weakenExpr (WCopy w) e3)
   EBuild1 x e1 e2 -> EBuild1 x (weakenExpr w e1) (weakenExpr (WCopy w) e2)
   EBuild x n es e -> EBuild x n (weakenVec w es) (weakenExpr (wcopyN n w) e)
   EFold1 x e1 e2 -> EFold1 x (weakenExpr (WCopy (WCopy w)) e1) (weakenExpr w e2)
@@ -211,6 +233,10 @@ weakenExpr w = \case
   EIdx x e1 es -> EIdx x (weakenExpr w e1) (weakenVec w es)
   EOp x op e -> EOp x op (weakenExpr w e)
   EMOne t i e -> EMOne t i (weakenExpr w e)
+  EMScope e -> EMScope (weakenExpr w e)
+  EMReturn t e -> EMReturn t (weakenExpr w e)
+  EMBind e1 e2 -> EMBind (weakenExpr w e1) (weakenExpr (WCopy w) e2)
+  EError t s -> EError t s
 
 wcopyN :: SNat n -> env :> env' -> ConsN n TIx env :> ConsN n TIx env'
 wcopyN SZ w = w
@@ -219,3 +245,7 @@ wcopyN (SS n) w = WCopy (wcopyN n w)
 weakenVec :: (env :> env') -> Vec n (Expr x env TIx) -> Vec n (Expr x env' TIx)
 weakenVec _ VNil = VNil
 weakenVec w (e :< v) = weakenExpr w e :< weakenVec w v
+
+slistMap :: (forall t. f t -> g t) -> SList f list -> SList g list
+slistMap _ SNil = SNil
+slistMap f (SCons x list) = SCons (f x) (slistMap f list)