1 files changed, 27 insertions, 1 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