From f6b5850f949eb671f0c7038db6dff80ca23ed946 Mon Sep 17 00:00:00 2001
From: Tom Smeding <tom@tomsmeding.com>
Date: Wed, 23 Jul 2025 21:09:59 +0200
Subject: WIP pruneExpr in AST.Count

---
 src/AST.hs | 29 ++++++++++++++++++++++++++++-
 1 file changed, 28 insertions(+), 1 deletion(-)

(limited to 'src/AST.hs')

diff --git a/src/AST.hs b/src/AST.hs
index 5aab4fc..b8bee1b 100644
--- a/src/AST.hs
+++ b/src/AST.hs
@@ -21,6 +21,7 @@ module AST (module AST, module AST.Types, module AST.Accum, module AST.Weaken) w
 
 import Data.Functor.Const
 import Data.Functor.Identity
+import Data.Int (Int64)
 import Data.Kind (Type)
 
 import Array
@@ -447,6 +448,16 @@ envKnown :: SList STy env -> Dict (KnownEnv env)
 envKnown SNil = Dict
 envKnown (t `SCons` env) | Dict <- styKnown t, Dict <- envKnown env = Dict
 
+cheapExpr :: Expr x env t -> Bool
+cheapExpr = \case
+  EVar{} -> True
+  ENil{} -> True
+  EConst{} -> True
+  EFst _ e -> cheapExpr e
+  ESnd _ e -> cheapExpr e
+  EUnit _ e -> cheapExpr e
+  _ -> False
+
 eTup :: SList (Ex env) list -> Ex env (Tup list)
 eTup = mkTup (ENil ext) (EPair ext)
 
@@ -516,6 +527,10 @@ eshapeEmpty (SS n) e =
                                       (EConst ext STI64 0)))
       (eshapeEmpty n (EFst ext (EVar ext (tTup (sreplicate (SS n) tIx)) IZ))))
 
+eshapeConst :: Shape n -> Ex env (Tup (Replicate n TIx))
+eshapeConst ShNil = ENil ext
+eshapeConst (sh `ShCons` n) = EPair ext (eshapeConst sh) (EConst ext STI64 (fromIntegral @Int @Int64 n))
+
 -- ezeroD2 :: STy t -> Ex env (ZeroInfo (D2 t)) -> Ex env (D2 t)
 -- ezeroD2 t ezi = EZero ext (d2M t) ezi
 
@@ -527,6 +542,7 @@ eshapeEmpty (SS n) e =
 
 eunPair :: Ex env (TPair a b) -> (forall env'. env :> env' -> Ex env' a -> Ex env' b -> Ex env' r) -> Ex env r
 eunPair (EPair _ e1 e2) k = k WId e1 e2
+eunPair e k | cheapExpr e = k WId (EFst ext e) (ESnd ext e)
 eunPair e k =
   elet e $
     k WSink
@@ -546,13 +562,24 @@ elet rhs body
   | Dict <- styKnown (typeOf rhs)
   = ELet ext rhs body
 
+-- | Let-bind it but don't use the value (just ensure the expression's effects don't get lost)
+use :: Ex env a -> Ex env b -> Ex env b
+use a b = elet a $ weakenExpr WSink b
+
 emaybe :: Ex env (TMaybe a) -> Ex env b -> (KnownTy a => Ex (a : env) b) -> Ex env b
 emaybe e a b
   | STMaybe t <- typeOf e
   , Dict <- styKnown t
   = EMaybe ext a b e
 
-elcase :: Ex env (TLEither a b) -> Ex env c -> (KnownTy a => Ex (a : env) c) -> (KnownTy b => Ex (b : env) c) -> Ex env c
+ecase :: Ex env (TEither a b) -> ((KnownTy a, KnownTy b) => Ex (a : env) c) -> ((KnownTy a, KnownTy b) => Ex (b : env) c) -> Ex env c
+ecase e a b
+  | STEither t1 t2 <- typeOf e
+  , Dict <- styKnown t1
+  , Dict <- styKnown t2
+  = ECase ext e a b
+
+elcase :: Ex env (TLEither a b) -> ((KnownTy a, KnownTy b) => Ex env c) -> ((KnownTy a, KnownTy b) => Ex (a : env) c) -> ((KnownTy a, KnownTy b) => Ex (b : env) c) -> Ex env c
 elcase e a b c
   | STLEither t1 t2 <- typeOf e
   , Dict <- styKnown t1
-- 
cgit v1.2.3-70-g09d2