From 9b7c3eea7e34f5eb0d91f93b803e853028c2cec8 Mon Sep 17 00:00:00 2001
From: Tom Smeding <tom@tomsmeding.com>
Date: Sat, 22 Nov 2025 22:41:09 +0100
Subject: WIP: Think about fusion

---
 src/CHAD/AST/Bindings.hs |  6 +++---
 src/CHAD/AST/Count.hs    | 12 ++++++------
 src/CHAD/AST/Pretty.hs   | 16 ++++++++--------
 3 files changed, 17 insertions(+), 17 deletions(-)

(limited to 'src/CHAD/AST')

diff --git a/src/CHAD/AST/Bindings.hs b/src/CHAD/AST/Bindings.hs
index c1a1e77..3ecda3e 100644
--- a/src/CHAD/AST/Bindings.hs
+++ b/src/CHAD/AST/Bindings.hs
@@ -28,7 +28,7 @@ data Bindings f env binds where
 deriving instance (forall e t. Show (f e t)) => Show (Bindings f env env')
 infixl `BPush`
 
-bpush :: Bindings (Expr x) env binds -> Expr x (Append binds env) t -> Bindings (Expr x) env (t : binds)
+bpush :: Bindings (Expr NoExt x) env binds -> Expr NoExt x (Append binds env) t -> Bindings (Expr NoExt x) env (t : binds)
 bpush b e = b `BPush` (typeOf e, e)
 infixl `bpush`
 
@@ -47,8 +47,8 @@ weakenBindings wf w (BPush b (t, x)) =
   in (BPush b' (t, wf w' x), WCopy w')
 
 weakenBindingsE :: env1 :> env2
-                -> Bindings (Expr x) env1 binds
-                -> (Bindings (Expr x) env2 binds, Append binds env1 :> Append binds env2)
+                -> Bindings (Expr NoExt x) env1 binds
+                -> (Bindings (Expr NoExt x) env2 binds, Append binds env1 :> Append binds env2)
 weakenBindingsE = weakenBindings weakenExpr
 
 weakenOver :: SList STy ts -> env :> env' -> Append ts env :> Append ts env'
diff --git a/src/CHAD/AST/Count.hs b/src/CHAD/AST/Count.hs
index 46173d2..1dad758 100644
--- a/src/CHAD/AST/Count.hs
+++ b/src/CHAD/AST/Count.hs
@@ -338,15 +338,15 @@ envMaskPrj (EMRest b) _ = b
 envMaskPrj (_ `EMPush` b) IZ = b
 envMaskPrj (env `EMPush` _) (IS i) = envMaskPrj env i
 
-occCount :: Idx env a -> Expr x env t -> Occ
+occCount :: Idx env a -> Expr NoExt x env t -> Occ
 occCount idx ex
   | Some env <- occCountAll ex
   = fst (occEnvPrj env idx)
 
-occCountAll :: Expr x env t -> Some (OccEnv Occ env)
+occCountAll :: Expr NoExt x env t -> Some (OccEnv Occ env)
 occCountAll ex = occCountX SsFull ex $ \env _ -> Some env
 
-pruneExpr :: SList f env -> Expr x env t -> Ex env t
+pruneExpr :: SList f env -> Expr NoExt x env t -> Ex env t
 pruneExpr env ex = occCountX SsFull ex $ \_ mkex -> mkex (fullOccEnv env)
   where
     fullOccEnv :: SList f env -> OccEnv () env env
@@ -365,7 +365,7 @@ pruneExpr env ex = occCountX SsFull ex $ \_ mkex -> mkex (fullOccEnv env)
 --   occurrence counts. The callback reconstructs a new expression in an
 --   updated "response" environment. The response must be at least as large as
 --   the computed usages.
-occCountX :: forall env t t' x r. Substruc t t' -> Expr x env t
+occCountX :: forall env t t' x r. Substruc t t' -> Expr NoExt x env t
           -> (forall env'. OccEnv Occ env env'
                            -- response OccEnv must be at least as large as the OccEnv returned above
                         -> (forall env''. OccEnv () env env'' -> Ex env'' t')
@@ -885,7 +885,7 @@ occCountX initialS topexpr k = case topexpr of
 
     handleReduction :: t ~ TArr n (TScal t2)
                     => (forall env2. Ex env2 (TArr (S n) (TScal t2)) -> Ex env2 (TArr n (TScal t2)))
-                    -> Expr x env (TArr (S n) (TScal t2))
+                    -> Expr NoExt x env (TArr (S n) (TScal t2))
                     -> r
     handleReduction reduce e
       | STArr (SS n) _ <- typeOf e =
@@ -914,7 +914,7 @@ deleteUnused (_ `SCons` env) (Some (OccPush occenv (Occ _ count) _)) k =
     case count of Zero -> k (SENo sub)
                   _    -> k (SEYesR sub)
 
-unsafeWeakenWithSubenv :: Subenv env env' -> Expr x env t -> Expr x env' t
+unsafeWeakenWithSubenv :: Subenv env env' -> Expr NoExt x env t -> Expr NoExt x env' t
 unsafeWeakenWithSubenv = \sub ->
   subst (\x t i -> case sinkViaSubenv i sub of
                      Just i' -> EVar x t i'
diff --git a/src/CHAD/AST/Pretty.hs b/src/CHAD/AST/Pretty.hs
index 9ddcb35..b763efe 100644
--- a/src/CHAD/AST/Pretty.hs
+++ b/src/CHAD/AST/Pretty.hs
@@ -63,20 +63,20 @@ nameBaseForType _ = "x"
 genName' :: String -> M String
 genName' prefix = (prefix ++) . show <$> genId
 
-genNameIfUsedIn' :: String -> STy a -> Idx env a -> Expr x env t -> M String
+genNameIfUsedIn' :: String -> STy a -> Idx env a -> Expr NoExt x env t -> M String
 genNameIfUsedIn' prefix ty idx ex
   | occCount idx ex == mempty = case ty of STNil -> return "()"
                                            _ -> return "_"
   | otherwise                 = genName' prefix
 
 -- TODO: let this return a type-tagged thing so that name environments are more typed than Const
-genNameIfUsedIn :: STy a -> Idx env a -> Expr x env t -> M String
+genNameIfUsedIn :: STy a -> Idx env a -> Expr NoExt x env t -> M String
 genNameIfUsedIn = \t -> genNameIfUsedIn' (nameBaseForType t) t
 
-pprintExpr :: (KnownEnv env, PrettyX x) => Expr x env t -> IO ()
+pprintExpr :: (KnownEnv env, PrettyX x) => Expr NoExt x env t -> IO ()
 pprintExpr = putStrLn . ppExpr knownEnv
 
-ppExpr :: PrettyX x => SList STy env -> Expr x env t -> String
+ppExpr :: PrettyX x => SList STy env -> Expr NoExt x env t -> String
 ppExpr senv e = render $ fst . flip runM 1 $ do
   val <- mkVal senv
   e' <- ppExpr' 0 val e
@@ -94,7 +94,7 @@ ppExpr senv e = render $ fst . flip runM 1 $ do
       name <- genName' "arg"
       return (Const name `SCons` val)
 
-ppExpr' :: PrettyX x => Int -> SVal env -> Expr x env t -> M ADoc
+ppExpr' :: PrettyX x => Int -> SVal env -> Expr NoExt x env t -> M ADoc
 ppExpr' d val expr = case expr of
   EVar _ _ i -> return $ ppString (getConst (slistIdx val i)) <> ppX expr
 
@@ -374,9 +374,9 @@ ppExpr' d val expr = case expr of
 
   EError _ _ s -> return $ ppParen (d > 10) $ ppString "error" <> ppX expr <+> ppString (show s)
 
-ppExprLet :: PrettyX x => Int -> SVal env -> Expr x env t -> M ADoc
+ppExprLet :: PrettyX x => Int -> SVal env -> Expr NoExt x env t -> M ADoc
 ppExprLet d val etop = do
-  let collect :: PrettyX x => SVal env -> Expr x env t -> M ([(String, Occ, ADoc)], ADoc)
+  let collect :: PrettyX x => SVal env -> Expr NoExt x env t -> M ([(String, Occ, ADoc)], ADoc)
       collect val' (ELet _ rhs body) = do
         let occ = occCount IZ body
         name <- genNameIfUsedIn (typeOf rhs) IZ body
@@ -426,7 +426,7 @@ ppCommut :: Commutative -> String
 ppCommut Commut = "(C)"
 ppCommut Noncommut = ""
 
-ppX :: PrettyX x => Expr x env t -> ADoc
+ppX :: PrettyX x => Expr NoExt x env t -> ADoc
 ppX expr = annotate AExt $ ppString $ prettyXsuffix (extOf expr)
 
 data Fixity = Prefix | Infix
-- 
cgit v1.2.3-70-g09d2