WIPsparse

1 files changed, 33 insertions, 2 deletions

diff --git a/src/AST/Sparse.hs b/src/AST/Sparse.hs
index ddae7fe..369d395 100644
--- a/src/AST/Sparse.hs
+++ b/src/AST/Sparse.hs
@@ -111,6 +111,37 @@ isAbsent (SpMaybe s) = isAbsent s
 isAbsent (SpArr s) = isAbsent s
 isAbsent SpScal = False
 
+sparsePlus :: SMTy t -> Sparse t t' -> Ex env t' -> Ex env t' -> Ex env t'
+sparsePlus _ SpAbsent _ _ = ENil ext
+sparsePlus t sp e1 e2 | Just Refl <- isDense t sp = EPlus ext t e1 e2
+sparsePlus t (SpSparse sp) e1 e2 = sparsePlus (SMTMaybe t) (SpMaybe sp) e1 e2  -- heh
+sparsePlus (SMTPair t1 t2) (SpPair sp1 sp2) e1 e2 =
+  eunPair e1 $ \w1 e1a e1b ->
+  eunPair (weakenExpr w1 e2) $ \w2 e2a e2b ->
+    EPair ext (sparsePlus t1 sp1 (weakenExpr w2 e1a) e2a)
+              (sparsePlus t2 sp2 (weakenExpr w2 e1b) e2b)
+sparsePlus (SMTLEither t1 t2) (SpLEither sp1 sp2) e1 e2 =
+  elet e2 $
+    elcase (weakenExpr WSink e1)
+      (evar IZ)
+      (elcase (evar (IS IZ))
+        (ELInl ext (applySparse sp2 (fromSMTy t2)) (evar IZ))
+        (ELInl ext (applySparse sp2 (fromSMTy t2)) (sparsePlus t1 sp1 (evar (IS IZ)) (evar IZ)))
+        (EError ext (fromSMTy (applySparse (SpLEither sp1 sp2) (SMTLEither t1 t2))) "splus ll+lr"))
+      (elcase (evar (IS IZ))
+        (ELInr ext (applySparse sp1 (fromSMTy t1)) (evar IZ))
+        (EError ext (fromSMTy (applySparse (SpLEither sp1 sp2) (SMTLEither t1 t2))) "splus lr+ll")
+        (ELInr ext (applySparse sp1 (fromSMTy t1)) (sparsePlus t2 sp2 (evar (IS IZ)) (evar IZ))))
+sparsePlus (SMTMaybe t) (SpMaybe sp) e1 e2 =
+  elet e2 $
+    emaybe (weakenExpr WSink e1)
+      (evar IZ)
+      (emaybe (evar (IS IZ))
+        (EJust ext (evar IZ))
+        (EJust ext (sparsePlus t sp (evar (IS IZ)) (evar IZ))))
+sparsePlus (SMTArr _ t) (SpArr sp) e1 e2 = ezipWith (sparsePlus t sp (evar (IS IZ)) (evar IZ)) e1 e2
+sparsePlus t@SMTScal{} SpScal e1 e2 = EPlus ext t e1 e2
+
 
 data SBool b where
   SF :: SBool False
@@ -120,7 +151,7 @@ deriving instance Show (SBool b)
 data Injection sp a b where
   -- | 'Inj' is purposefully also allowed when @sp@ is @False@ so that
   -- 'sparsePlusS' can provide injections even if the caller doesn't require
-  -- them. This eliminates pointless checks.
+  -- them. This simplifies the sparsePlusS code.
   Inj :: (forall e. Ex e a -> Ex e b) -> Injection sp a b
   Noinj :: Injection False a b
 
@@ -138,7 +169,7 @@ withInj2 Noinj _ _ = Noinj
 withInj2 _ Noinj _ = Noinj
 
 -- | This function produces quadratically-sized code in the presence of nested
--- dynamic sparsity. しょうがない。
+-- dynamic sparsity. TODO can this be improved?
 sparsePlusS
   :: SBool inj1 -> SBool inj2
   -> SMTy t -> Sparse t t1 -> Sparse t t2