summaryrefslogtreecommitdiff
path: root/src/CHAD.hs
diff options
context:
space:
mode:
Diffstat (limited to 'src/CHAD.hs')
-rw-r--r--src/CHAD.hs101
1 files changed, 79 insertions, 22 deletions
diff --git a/src/CHAD.hs b/src/CHAD.hs
index 59d61a7..8b9f17a 100644
--- a/src/CHAD.hs
+++ b/src/CHAD.hs
@@ -28,7 +28,6 @@ module CHAD (
Select,
) where
-import Data.Bifunctor (first, second)
import Data.Functor.Const
import Data.Kind (Type)
import GHC.Stack (HasCallStack)
@@ -198,6 +197,7 @@ 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
@@ -224,11 +224,13 @@ subDescr (des `DPush` (t, sto)) (SEYes sub) k =
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
@@ -240,8 +242,13 @@ 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)
---------------------------------- DERIVATIVES ---------------------------------
@@ -338,13 +345,27 @@ conv1Idx :: Idx env t -> Idx (D1E env) (D1 t)
conv1Idx IZ = IZ
conv1Idx (IS i) = IS (conv1Idx i)
-conv2Idx :: Descr env sto -> Idx env t
- -> Either (Idx (D2AcE (Select env sto "accum")) (TAccum (D2 t)))
- (Idx (Select env sto "merge") t)
-conv2Idx (DPush _ (_, SAccum)) IZ = Left IZ
-conv2Idx (DPush _ (_, SMerge)) IZ = Right IZ
-conv2Idx (DPush des (_, SAccum)) (IS i) = first IS (conv2Idx des i)
-conv2Idx (DPush des (_, SMerge)) (IS i) = second IS (conv2Idx des i)
+data Idx2 env sto t
+ = Idx2Ac (Idx (D2AcE (Select env sto "accum")) (TAccum (D2 t)))
+ | Idx2Me (Idx (Select env sto "merge") t)
+ | Idx2Di (Idx (Select env sto "discr") t)
+
+conv2Idx :: Descr env sto -> Idx env t -> Idx2 env sto t
+conv2Idx (DPush _ (_, SAccum)) IZ = Idx2Ac IZ
+conv2Idx (DPush _ (_, SMerge)) IZ = Idx2Me IZ
+conv2Idx (DPush _ (_, SDiscr)) IZ = Idx2Di IZ
+conv2Idx (DPush des (_, SAccum)) (IS i) =
+ case conv2Idx des i of Idx2Ac j -> Idx2Ac (IS j)
+ Idx2Me j -> Idx2Me j
+ Idx2Di j -> Idx2Di j
+conv2Idx (DPush des (_, SMerge)) (IS i) =
+ case conv2Idx des i of Idx2Ac j -> Idx2Ac j
+ Idx2Me j -> Idx2Me (IS j)
+ Idx2Di j -> Idx2Di j
+conv2Idx (DPush des (_, SDiscr)) (IS i) =
+ case conv2Idx des i of Idx2Ac j -> Idx2Ac j
+ Idx2Me j -> Idx2Me j
+ Idx2Di j -> Idx2Di (IS j)
conv2Idx DTop i = case i of {}
@@ -536,6 +557,13 @@ accumPromote pdty (descr `DPush` (t :: STy t, sto)) k =
(#acc :++: (#d :++: #shb :++: #tl)))
SMerge -> case t of
+ -- Discrete values are left as-is
+ _ | isDiscrete t ->
+ k (storepl `DPush` (t, SDiscr))
+ envpro
+ (SENo prosub)
+ wf
+
-- Arrays with "merge" storage are promoted to an accumulator in envPro
STArr (arrn :: SNat arrn) (arrt :: STy arrt) ->
k (storepl `DPush` (t, SAccum))
@@ -555,18 +583,40 @@ accumPromote pdty (descr `DPush` (t :: STy t, sto)) k =
.> WPick @(TAccum (D2 (TArr arrn arrt))) @(D2 dt : shbinds) (Const () `SCons` shbindsC)
(WId @(D2AcE (Select env1 stoRepl "accum"))))
- -- "merge" values must be an array, so reject everything else. (TODO: generalise this)
+ -- "merge" values must be an array or fully discrete, so reject everything else. (TODO: generalise this)
_ ->
- error $ "Closure variable of 'build'-like thing contains a non-array SMerge value: " ++ show t
- -- where
- -- containsTArr :: STy t' -> Bool
- -- containsTArr = \case
- -- STNil -> False
- -- STPair a b -> containsTArr a || containsTArr b
- -- STEither a b -> containsTArr a || containsTArr b
- -- STArr{} -> True
- -- STScal{} -> False
- -- STAccum{} -> error "An accumulator in merge storage?"
+ error $ "Closure variable of 'build'-like thing contains a non-array non-discrete SMerge value: " ++ show t
+
+ -- Discrete values are left as-is, nothing to do
+ SDiscr ->
+ k (storepl `DPush` (t, SDiscr))
+ envpro
+ prosub
+ wf
+ where
+ isDiscrete :: STy t' -> Bool
+ isDiscrete = \case
+ STNil -> True
+ STPair a b -> isDiscrete a && isDiscrete b
+ STEither a b -> isDiscrete a && isDiscrete b
+ STMaybe a -> isDiscrete a
+ STArr _ a -> isDiscrete a
+ STScal st -> case st of
+ STI32 -> True
+ STI64 -> True
+ STF32 -> False
+ STF64 -> False
+ STBool -> True
+ STAccum{} -> False
+
+ -- containsTArr :: STy t' -> Bool
+ -- containsTArr = \case
+ -- STNil -> False
+ -- STPair a b -> containsTArr a || containsTArr b
+ -- STEither a b -> containsTArr a || containsTArr b
+ -- STArr{} -> True
+ -- STScal{} -> False
+ -- STAccum{} -> error "An accumulator in merge storage?"
makeAccumulators :: SList STy envPro -> Ex (Append (D2AcE envPro) env) t -> Ex env (InvTup t (D2E envPro))
makeAccumulators SNil e = e
@@ -682,20 +732,27 @@ drev :: forall env sto t.
drev des = \case
EVar _ t i ->
case conv2Idx des i of
- Left accI ->
+ Idx2Ac accI ->
Ret BTop
SETop
(EVar ext (d1 t) (conv1Idx i))
(subenvNone (select SMerge des))
(EAccum SZ (ENil ext) (EVar ext (d2 t) IZ) (EVar ext (STAccum (d2 t)) (IS accI)))
- Right tupI ->
+ Idx2Me tupI ->
Ret BTop
SETop
(EVar ext (d1 t) (conv1Idx i))
(subenvOnehot (select SMerge des) tupI)
(EPair ext (ENil ext) (EVar ext (d2 t) IZ))
+ Idx2Di _ ->
+ Ret BTop
+ SETop
+ (EVar ext (d1 t) (conv1Idx i))
+ (subenvNone (select SMerge des))
+ (ENil ext)
+
ELet _ rhs body
| Ret (rhs0 :: Bindings _ _ rhs_shbinds) (subtapeRHS :: Subenv _ rhs_tapebinds) (rhs1 :: Ex _ d1_a) subRHS rhs2
<- drev des rhs
@@ -929,7 +986,7 @@ drev des = \case
let e = unsafeWeakenWithSubenv (SEYes usedSub) orige in
subDescr des usedSub $ \usedDes subMergeUsed subAccumUsed subD1eUsed ->
accumPromote eltty usedDes $ \prodes (envPro :: SList _ envPro) proSub wPro ->
- case drev (prodes `DPush` (shty, SMerge)) e of { Ret (e0 :: Bindings _ _ e_binds) (subtapeE :: Subenv _ e_tape) e1 sub e2 ->
+ case drev (prodes `DPush` (shty, SDiscr)) e of { Ret (e0 :: Bindings _ _ e_binds) (subtapeE :: Subenv _ e_tape) e1 sub e2 ->
case assertSubenvEmpty sub of { Refl ->
let tapety = tapeTy (subList (bindingsBinds e0) subtapeE) in
let collectexpr = bindingsCollect e0 subtapeE in
generated by cgit v1.2.3-70-g09d2 (git 2.48.1) at 2025-03-31 23:29:26 +0200