WIP accum top-level args

1 files changed, 75 insertions, 0 deletions

diff --git a/src/CHAD/EnvDescr.hs b/src/CHAD/EnvDescr.hs
new file mode 100644
index 0000000..fcd91f7
--- /dev/null
+++ b/src/CHAD/EnvDescr.hs
@@ -0,0 +1,75 @@
+{-# LANGUAGE DataKinds #-}
+{-# LANGUAGE GADTs #-}
+{-# LANGUAGE PolyKinds #-}
+{-# LANGUAGE RankNTypes #-}
+{-# LANGUAGE StandaloneDeriving #-}
+{-# LANGUAGE StandaloneKindSignatures #-}
+{-# LANGUAGE TypeFamilies #-}
+{-# LANGUAGE TypeOperators #-}
+module CHAD.EnvDescr where
+
+import Data.Kind (Type)
+import GHC.TypeLits (Symbol)
+
+import AST.Env
+import AST.Types
+import CHAD.Types
+import Data
+
+
+type Storage :: Symbol -> Type
+data Storage s where
+  SAccum :: Storage "accum"  -- ^ in the monad state as a mutable accumulator
+  SMerge :: Storage "merge"  -- ^ just return and merge
+  SDiscr :: Storage "discr"  -- ^ we happen to know this is a discrete type and won't need any contributions
+deriving instance Show (Storage s)
+
+-- | Environment description
+data Descr env sto where
+  DTop :: Descr '[] '[]
+  DPush :: Descr env sto -> (STy t, Storage s) -> Descr (t : env) (s : sto)
+deriving instance Show (Descr env sto)
+
+descrList :: Descr env sto -> SList STy env
+descrList DTop = SNil
+descrList (des `DPush` (t, _)) = t `SCons` descrList des
+
+-- | This could have more precise typing on the output storage.
+subDescr :: Descr env sto -> Subenv env env'
+         -> (forall sto'. Descr env' sto'
+                       -> Subenv (Select env sto "merge") (Select env' sto' "merge")
+                       -> Subenv (D2AcE (Select env sto "accum")) (D2AcE (Select env' sto' "accum"))
+                       -> Subenv (D1E env) (D1E env')
+                       -> r)
+         -> r
+subDescr DTop SETop k = k DTop SETop SETop SETop
+subDescr (des `DPush` (t, sto)) (SEYes sub) k =
+  subDescr des sub $ \des' submerge subaccum subd1e ->
+    case sto of
+      SMerge -> k (des' `DPush` (t, sto)) (SEYes submerge) subaccum (SEYes subd1e)
+      SAccum -> k (des' `DPush` (t, sto)) submerge (SEYes subaccum) (SEYes subd1e)
+      SDiscr -> k (des' `DPush` (t, sto)) submerge subaccum (SEYes subd1e)
+subDescr (des `DPush` (_, sto)) (SENo sub) k =
+  subDescr des sub $ \des' submerge subaccum subd1e ->
+    case sto of
+      SMerge -> k des' (SENo submerge) subaccum (SENo subd1e)
+      SAccum -> k des' submerge (SENo subaccum) (SENo subd1e)
+      SDiscr -> k des' submerge subaccum (SENo subd1e)
+
+-- | Select only the types from the environment that have the specified storage
+type family Select env sto s where
+  Select '[] '[] _ = '[]
+  Select (t : ts) (s : sto) s = t : Select ts sto s
+  Select (_ : ts) (_ : sto) s = Select ts sto s
+
+select :: Storage s -> Descr env sto -> SList STy (Select env sto s)
+select _ DTop = SNil
+select s@SAccum (DPush des (t, SAccum)) = SCons t (select s des)
+select s@SMerge (DPush des (_, SAccum)) = select s des
+select s@SDiscr (DPush des (_, SAccum)) = select s des
+select s@SAccum (DPush des (_, SMerge)) = select s des
+select s@SMerge (DPush des (t, SMerge)) = SCons t (select s des)
+select s@SDiscr (DPush des (_, SMerge)) = select s des
+select s@SAccum (DPush des (_, SDiscr)) = select s des
+select s@SMerge (DPush des (_, SDiscr)) = select s des
+select s@SDiscr (DPush des (t, SDiscr)) = SCons t (select s des)