WIP Build1

1 files changed, 12 insertions, 0 deletions

diff --git a/src/AST.hs b/src/AST.hs
index d9acd99..c191651 100644
--- a/src/AST.hs
+++ b/src/AST.hs
@@ -15,6 +15,7 @@
 {-# LANGUAGE EmptyCase #-}
 module AST (module AST, module AST.Weaken) where
 
+import Data.Bifunctor (first)
 import Data.Functor.Const
 import Data.Kind (Type)
 import Data.Int
@@ -55,6 +56,9 @@ deriving instance Show (SScalTy t)
 
 type TIx = TScal TI64
 
+tIx :: STy TIx
+tIx = STScal STI64
+
 type family ScalRep t where
   ScalRep TI32 = Int32
   ScalRep TI64 = Int64
@@ -92,6 +96,7 @@ data Expr x env t where
 
   -- expression operations
   EConst :: Show (ScalRep t) => x (TScal t) -> SScalTy t -> ScalRep t -> Expr x env (TScal t)
+  EIdx0 :: x t -> Expr x env (TArr Z t) -> Expr x env t
   EIdx1 :: x (TArr n t) -> Expr x env (TArr (S n) t) -> Expr x env TIx -> Expr x env (TArr n t)
   EIdx :: x t -> Expr x env (TArr n t) -> Vec n (Expr x env TIx) -> Expr x env t
   EOp :: x t -> SOp a t -> Expr x env a -> Expr x env t
@@ -150,6 +155,7 @@ typeOf = \case
   EFold1 _ _ e | STArr (SS n) t <- typeOf e -> STArr n t
 
   EConst _ t _ -> STScal t
+  EIdx0 _ e | STArr _ t <- typeOf e -> t
   EIdx1 _ e _ | STArr (SS n) t <- typeOf e -> STArr n t
   EIdx _ e _ | STArr _ t <- typeOf e -> t
   EOp _ op _ -> opt2 op
@@ -210,6 +216,7 @@ subst' f w = \case
   EBuild x es e -> EBuild x (fmap (subst' f w) es) (subst' (sinkFN (vecLength es) f) (wcopyN (vecLength es) w) e)
   EFold1 x a b -> EFold1 x (subst' (sinkF (sinkF f)) (WCopy (WCopy w)) a) (subst' f w b)
   EConst x t v -> EConst x t v
+  EIdx0 x e -> EIdx0 x (subst' f w e)
   EIdx1 x a b -> EIdx1 x (subst' f w a) (subst' f w b)
   EIdx x e es -> EIdx x (subst' f w e) (fmap (subst' f w) es)
   EOp x op e -> EOp x op (subst' f w e)
@@ -254,6 +261,11 @@ idx2int :: Idx env t -> Int
 idx2int IZ = 0
 idx2int (IS n) = 1 + idx2int n
 
+splitIdx :: forall env2 env1 t f. SList f env1 -> Idx (Append env1 env2) t -> Either (Idx env1 t) (Idx env2 t)
+splitIdx SNil i = Right i
+splitIdx (SCons _ _) IZ = Left IZ
+splitIdx (SCons _ l) (IS i) = first IS (splitIdx l i)
+
 class KnownScalTy t where knownScalTy :: SScalTy t
 instance KnownScalTy TI32 where knownScalTy = STI32
 instance KnownScalTy TI64 where knownScalTy = STI64