WIP

4 files changed, 84 insertions, 9 deletions

diff --git a/src/AST/Count.hs b/src/AST/Count.hs
index 289c1fb..a4ff9f2 100644
--- a/src/AST/Count.hs
+++ b/src/AST/Count.hs
@@ -2,12 +2,14 @@
 {-# LANGUAGE DeriveGeneric #-}
 {-# LANGUAGE DerivingStrategies #-}
 {-# LANGUAGE DerivingVia #-}
+{-# LANGUAGE EmptyCase #-}
 {-# LANGUAGE GADTs #-}
 {-# LANGUAGE LambdaCase #-}
 {-# LANGUAGE PolyKinds #-}
 {-# LANGUAGE QuantifiedConstraints #-}
 {-# LANGUAGE RankNTypes #-}
 {-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE StandaloneDeriving #-}
 {-# LANGUAGE TypeOperators #-}
 module AST.Count where
 
@@ -15,6 +17,7 @@ import Data.Functor.Const
 import GHC.Generics (Generic, Generically(..))
 
 import AST
+import AST.Env
 import Data
 
 
@@ -110,9 +113,10 @@ occCountGeneral onehot unpush unpushN alter many = go
       EInr _ _ e -> go e
       ECase _ e a b -> go e <> (unpush (go a) `alter` unpush (go b))
       EBuild1 _ a b -> go a <> many (unpush (go b))
-      EBuild _ es e -> foldMap go es <> many (unpushN (vecLength es) (go e))
+      EBuild _ n a b -> go a <> many (unpushN n (go b))
       EFold1 _ a b -> many (unpush (unpush (go a))) <> go b
       EUnit _ e -> go e
+      EReplicate _ e -> go e
       EConst{} -> mempty
       EIdx0 _ e -> go e
       EIdx1 _ a b -> go a <> go b
@@ -121,3 +125,25 @@ occCountGeneral onehot unpush unpushN alter many = go
       EWith a b -> go a <> unpush (go b)
       EAccum1 a b e -> go a <> go b <> go e
       EError{} -> mempty
+
+
+deleteUnused :: SList f env -> OccEnv env -> (forall env'. Subenv env env' -> r) -> r
+deleteUnused SNil OccEnd k = k SETop
+deleteUnused (_ `SCons` env) OccEnd k =
+  deleteUnused env OccEnd $ \sub -> k (SENo sub)
+deleteUnused (_ `SCons` env) (OccPush occenv (Occ _ count)) k =
+  deleteUnused env occenv $ \sub ->
+    case count of Zero -> k (SENo sub)
+                  _    -> k (SEYes sub)
+
+unsafeWeakenWithSubenv :: Subenv env env' -> Expr x env t -> Expr x env' t
+unsafeWeakenWithSubenv = \sub ->
+  subst (\x t i -> case sinkWithSubenv i sub of
+                     Just i' -> EVar x t i'
+                     Nothing -> error "unsafeWeakenWithSubenv: Index occurred that was subenv'd away")
+  where
+    sinkWithSubenv :: Idx env t -> Subenv env env' -> Maybe (Idx env' t)
+    sinkWithSubenv IZ (SEYes _) = Just IZ
+    sinkWithSubenv IZ (SENo _) = Nothing
+    sinkWithSubenv (IS i) (SEYes sub) = IS <$> sinkWithSubenv i sub
+    sinkWithSubenv (IS i) (SENo sub) = sinkWithSubenv i sub
diff --git a/src/AST/Env.hs b/src/AST/Env.hs
new file mode 100644
index 0000000..c33bad3
--- /dev/null
+++ b/src/AST/Env.hs
@@ -0,0 +1,43 @@
+{-# LANGUAGE DataKinds #-}
+{-# LANGUAGE EmptyCase #-}
+{-# LANGUAGE GADTs #-}
+{-# LANGUAGE PolyKinds #-}
+{-# LANGUAGE StandaloneDeriving #-}
+{-# LANGUAGE TypeOperators #-}
+module AST.Env where
+
+import AST.Weaken
+import Data
+
+
+-- | @env'@ is a subset of @env@: each element of @env@ is either included in
+-- @env'@ ('SEYes') or not included in @env'@ ('SENo').
+data Subenv env env' where
+  SETop :: Subenv '[] '[]
+  SEYes :: Subenv env env' -> Subenv (t : env) (t : env')
+  SENo  :: Subenv env env' -> Subenv (t : env) env'
+deriving instance Show (Subenv env env')
+
+subList :: SList f env -> Subenv env env' -> SList f env'
+subList SNil SETop = SNil
+subList (SCons x xs) (SEYes sub) = SCons x (subList xs sub)
+subList (SCons _ xs) (SENo sub) = subList xs sub
+
+subenvAll :: SList f env -> Subenv env env
+subenvAll SNil = SETop
+subenvAll (SCons _ env) = SEYes (subenvAll env)
+
+subenvNone :: SList f env -> Subenv env '[]
+subenvNone SNil = SETop
+subenvNone (SCons _ env) = SENo (subenvNone env)
+
+subenvOnehot :: SList f env -> Idx env t -> Subenv env '[t]
+subenvOnehot (SCons _ env) IZ = SEYes (subenvNone env)
+subenvOnehot (SCons _ env) (IS i) = SENo (subenvOnehot env i)
+subenvOnehot SNil i = case i of {}
+
+subenvCompose :: Subenv env1 env2 -> Subenv env2 env3 -> Subenv env1 env3
+subenvCompose SETop SETop = SETop
+subenvCompose (SEYes sub1) (SEYes sub2) = SEYes (subenvCompose sub1 sub2)
+subenvCompose (SEYes sub1) (SENo sub2) = SENo (subenvCompose sub1 sub2)
+subenvCompose (SENo sub1) sub2 = SENo (subenvCompose sub1 sub2)
diff --git a/src/AST/Pretty.hs b/src/AST/Pretty.hs
index 1dc9dd3..dbbc021 100644
--- a/src/AST/Pretty.hs
+++ b/src/AST/Pretty.hs
@@ -113,14 +113,12 @@ ppExpr' d val = \case
     return $ showParen (d > 10) $
       showString "build1 " . a' . showString (" (\\" ++ name ++ " -> ") . b' . showString ")"
 
-  EBuild _ es e -> do
-    es' <- mapM (ppExpr' 0 val) es
-    names <- mapM (const genName) es  -- TODO generate underscores
-    e' <- ppExpr' 0 (vpushN names val) e
+  EBuild _ n a b -> do
+    a' <- ppExpr' 11 val a
+    names <- sequence (vecGenerate n (\_ -> genName))  -- TODO generate underscores
+    e' <- ppExpr' 0 (vpushN names val) b
     return $ showParen (d > 10) $
-      showString "build ["
-      . foldr (.) id (intersperse (showString ", ") (reverse (toList es')))
-      . showString "] (\\["
+      showString "build " . a' . showString " (\\["
       . foldr (.) id (intersperse (showString ",") (map showString (reverse (toList names))))
       . showString ("] -> ") . e' . showString ")"
 
@@ -137,6 +135,10 @@ ppExpr' d val = \case
     e' <- ppExpr' 11 val e
     return $ showParen (d > 10) $ showString "unit " . e'
 
+  EReplicate _ e -> do
+    e' <- ppExpr' 11 val e
+    return $ showParen (d > 10) $ showString "replicate " . e'
+
   EConst _ ty v -> return $ showString $ case ty of
     STI32 -> show v ; STI64 -> show v ; STF32 -> show v ; STF64 -> show v ; STBool -> show v
 
diff --git a/src/AST/Weaken.hs b/src/AST/Weaken.hs
index e0b5232..78276ca 100644
--- a/src/AST/Weaken.hs
+++ b/src/AST/Weaken.hs
@@ -39,7 +39,7 @@ splitIdx (SCons _ l) (IS i) = first IS (splitIdx l i)
 data env :> env' where
   WId :: env :> env
   WSink :: forall t env. env :> (t : env)
-  WCopy :: env :> env' -> (t : env) :> (t : env')
+  WCopy :: forall t env env'. env :> env' -> (t : env) :> (t : env')
   WPop :: (t : env) :> env' -> env :> env'
   WThen :: env1 :> env2 -> env2 :> env3 -> env1 :> env3
   WClosed :: SList (Const ()) env -> '[] :> env
@@ -95,6 +95,10 @@ wSinks :: forall env bs f. SList f bs -> env :> Append bs env
 wSinks SNil = WId
 wSinks (SCons _ spine) = WSink .> wSinks spine
 
+wSinksAnd :: forall env env' bs f. SList f bs -> env :> env' -> env :> Append bs env'
+wSinksAnd SNil w = w
+wSinksAnd (SCons _ spine) w = WSink .> wSinksAnd spine w
+
 wCopies :: SList f bs -> env1 :> env2 -> Append bs env1 :> Append bs env2
 wCopies SNil w = w
 wCopies (SCons _ spine) w = WCopy (wCopies spine w)