{-# LANGUAGE GADTs #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE ScopedTypeVariables #-}

{-# OPTIONS_GHC -fmax-pmcheck-models=60 #-}
module AST.Sparse where

import Data.Kind (Constraint, Type)
import Data.Type.Equality

import AST


data Sparse t t' where
  SpDense :: Sparse t t
  SpSparse :: Sparse t t' -> Sparse t (TMaybe t')
  SpAbsent :: Sparse t TNil

  SpPair :: Sparse a a' -> Sparse b b' -> Sparse (TPair a b) (TPair a' b')
  SpLEither :: Sparse a a' -> Sparse b b' -> Sparse (TLEither a b) (TLEither a' b')
  SpLeft :: Sparse a a' -> Sparse (TLEither a b) a'
  SpRight :: Sparse b b' -> Sparse (TLEither a b) b'
  SpMaybe :: Sparse t t' -> Sparse (TMaybe t) (TMaybe t')
  SpJust :: Sparse t t' -> Sparse (TMaybe t) t'
  SpArr :: Sparse t t' -> Sparse (TArr n t) (TArr n t')
deriving instance Show (Sparse t t')

applySparse :: Sparse t t' -> STy t -> STy t'
applySparse SpDense t = t
applySparse (SpSparse s) t = STMaybe (applySparse s t)
applySparse SpAbsent _ = STNil
applySparse (SpPair s1 s2) (STPair t1 t2) = STPair (applySparse s1 t1) (applySparse s2 t2)
applySparse (SpLEither s1 s2) (STLEither t1 t2) = STLEither (applySparse s1 t1) (applySparse s2 t2)
applySparse (SpLeft s) (STLEither t1 _) = applySparse s t1
applySparse (SpRight s) (STLEither _ t2) = applySparse s t2
applySparse (SpMaybe s) (STMaybe t) = STMaybe (applySparse s t)
applySparse (SpJust s) (STMaybe t) = applySparse s t
applySparse (SpArr s) (STArr n t) = STArr n (applySparse s t)


class IsSubType s where
  type IsSubTypeSubject s (f :: k -> Type) :: Constraint
  subtApply :: IsSubTypeSubject s f => s t t' -> f t -> f t'
  subtTrans :: s a b -> s b c -> s a c
  subtFull :: s a a

instance IsSubType (:~:) where
  type IsSubTypeSubject (:~:) f = ()
  subtApply = gcastWith
  subtTrans = trans
  subtFull = Refl

instance IsSubType Sparse where
  type IsSubTypeSubject Sparse f = f ~ STy
  subtApply = applySparse

  subtTrans SpDense s = s
  subtTrans s SpDense = s
  subtTrans s1 (SpSparse s2) = SpSparse (subtTrans s1 s2)
  subtTrans _ SpAbsent = SpAbsent
  subtTrans (SpPair s1a s1b) (SpPair s2a s2b) = SpPair (subtTrans s1a s2a) (subtTrans s1b s2b)
  subtTrans (SpLEither s1a s1b) (SpLEither s2a s2b) = SpLEither (subtTrans s1a s2a) (subtTrans s1b s2b)
  subtTrans (SpLEither s1 _) (SpLeft s2) = SpLeft (subtTrans s1 s2)
  subtTrans (SpLEither _ s1) (SpRight s2) = SpRight (subtTrans s1 s2)
  subtTrans (SpLeft s1) s2 = SpLeft (subtTrans s1 s2)
  subtTrans (SpRight s1) s2 = SpRight (subtTrans s1 s2)
  subtTrans (SpSparse s1) (SpMaybe s2) = SpSparse (subtTrans s1 s2)
  subtTrans (SpSparse s1) (SpJust s2) = subtTrans s1 s2
  subtTrans (SpMaybe s1) (SpMaybe s2) = SpMaybe (subtTrans s1 s2)
  subtTrans (SpMaybe s1) (SpJust s2) = SpJust (subtTrans s1 s2)
  subtTrans (SpJust s1) s2 = SpJust (subtTrans s1 s2)
  subtTrans (SpArr s1) (SpArr s2) = SpArr (subtTrans s1 s2)

  subtFull = SpDense


data SBool b where
  SF :: SBool False
  ST :: SBool True
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.
  Inj :: (forall e. Ex e a -> Ex e b) -> Injection sp a b
  Noinj :: Injection False a b

withInj :: Injection sp a b -> ((forall e. Ex e a -> Ex e b) -> (forall e'. Ex e' a' -> Ex e' b')) -> Injection sp a' b'
withInj (Inj f) k = Inj (k f)
withInj Noinj _ = Noinj

withInj2 :: Injection sp a1 b1 -> Injection sp a2 b2
         -> ((forall e. Ex e a1 -> Ex e b1)
             -> (forall e. Ex e a2 -> Ex e b2)
             -> (forall e'. Ex e' a' -> Ex e' b'))
         -> Injection sp a' b'
withInj2 (Inj f) (Inj g) k = Inj (k f g)
withInj2 Noinj _ _ = Noinj
withInj2 _ Noinj _ = Noinj

-- | This function produces quadratically-sized code in the presence of nested
-- dynamic sparsity. しょうがない。
sparsePlusS
  :: SBool inj1 -> SBool inj2
  -> SMTy t -> Sparse t t1 -> Sparse t t2
  -> (forall t3. Sparse t t3
              -> Injection inj1 t1 t3  -- only available if first injection is requested (second argument may be absent)
              -> Injection inj2 t2 t3  -- only available if second injection is requested (first argument may be absent)
              -> (forall e. Ex e t1 -> Ex e t2 -> Ex e t3)
              -> r)
  -> r
-- nil override
sparsePlusS _ _ SMTNil _ _ k = k SpAbsent (Inj $ \_ -> ENil ext) (Inj $ \_ -> ENil ext) (\_ _ -> ENil ext)

-- simplifications
sparsePlusS req1 req2 t (SpSparse SpAbsent) sp2 k =
  sparsePlusS req1 req2 t SpAbsent sp2 $ \sp3 minj1 minj2 plus ->
    k sp3 (withInj minj1 $ \inj1 -> \_ -> inj1 (ENil ext)) minj2 (\_ b -> plus (ENil ext) b)
sparsePlusS req1 req2 t sp1 (SpSparse SpAbsent) k =
  sparsePlusS req1 req2 t sp1 SpAbsent $ \sp3 minj1 minj2 plus ->
    k sp3 minj1 (withInj minj2 $ \inj2 -> \_ -> inj2 (ENil ext)) (\a _ -> plus a (ENil ext))

sparsePlusS req1 req2 t (SpSparse (SpSparse sp1)) sp2 k =
  let ta = applySparse sp1 (fromSMTy t) in
  sparsePlusS req1 req2 t (SpSparse sp1) sp2 $ \sp3 minj1 minj2 plus ->
    k sp3
      (withInj minj1 $ \inj1 -> \a -> inj1 (emaybe a (ENothing ext ta) (EVar ext (STMaybe ta) IZ)))
      minj2
      (\a b -> plus (emaybe a (ENothing ext ta) (EVar ext (STMaybe ta) IZ)) b)
sparsePlusS req1 req2 t sp1 (SpSparse (SpSparse sp2)) k =
  let tb = applySparse sp2 (fromSMTy t) in
  sparsePlusS req1 req2 t sp1 (SpSparse sp2) $ \sp3 minj1 minj2 plus ->
    k sp3
      minj1
      (withInj minj2 $ \inj2 -> \b -> inj2 (emaybe b (ENothing ext tb) (EVar ext (STMaybe tb) IZ)))
      (\a b -> plus a (emaybe b (ENothing ext tb) (EVar ext (STMaybe tb) IZ)))

sparsePlusS req1 req2 t (SpSparse (SpLEither sp1a sp1b)) sp2 k =
  let STLEither ta tb = applySparse (SpLEither sp1a sp1b) (fromSMTy t) in
  sparsePlusS req1 req2 t (SpLEither sp1a sp1b) sp2 $ \sp3 minj1 minj2 plus ->
    k sp3
      (withInj minj1 $ \inj1 -> \a -> inj1 (emaybe a (ELNil ext ta tb) (EVar ext (STLEither ta tb) IZ)))
      minj2
      (\a b -> plus (emaybe a (ELNil ext ta tb) (EVar ext (STLEither ta tb) IZ)) b)
sparsePlusS req1 req2 t sp1 (SpSparse (SpLEither sp2a sp2b)) k =
  let STLEither ta tb = applySparse (SpLEither sp2a sp2b) (fromSMTy t) in
  sparsePlusS req1 req2 t sp1 (SpLEither sp2a sp2b) $ \sp3 minj1 minj2 plus ->
    k sp3
      minj1
      (withInj minj2 $ \inj2 -> \b -> inj2 (emaybe b (ELNil ext ta tb) (EVar ext (STLEither ta tb) IZ)))
      (\a b -> plus a (emaybe b (ELNil ext ta tb) (EVar ext (STLEither ta tb) IZ)))

sparsePlusS req1 req2 t (SpSparse (SpMaybe sp1)) sp2 k =
  let STMaybe ta = applySparse (SpMaybe sp1) (fromSMTy t) in
  sparsePlusS req1 req2 t (SpMaybe sp1) sp2 $ \sp3 minj1 minj2 plus ->
    k sp3
      (withInj minj1 $ \inj1 -> \a -> inj1 (emaybe a (ENothing ext ta) (evar IZ)))
      minj2
      (\a b -> plus (emaybe a (ENothing ext ta) (EVar ext (STMaybe ta) IZ)) b)
sparsePlusS req1 req2 t sp1 (SpSparse (SpMaybe sp2)) k =
  let STMaybe tb = applySparse (SpMaybe sp2) (fromSMTy t) in
  sparsePlusS req1 req2 t sp1 (SpMaybe sp2) $ \sp3 minj1 minj2 plus ->
    k sp3
      minj1
      (withInj minj2 $ \inj2 -> \b -> inj2 (emaybe b (ENothing ext tb) (evar IZ)))
      (\a b -> plus a (emaybe b (ENothing ext tb) (EVar ext (STMaybe tb) IZ)))
sparsePlusS req1 req2 t (SpMaybe (SpSparse sp1)) sp2 k = sparsePlusS req1 req2 t (SpSparse (SpMaybe sp1)) sp2 k
sparsePlusS req1 req2 t sp1 (SpMaybe (SpSparse sp2)) k = sparsePlusS req1 req2 t sp1 (SpSparse (SpMaybe sp2)) k

-- TODO: sparse of Just is just Maybe

-- dense plus
sparsePlusS _ _ t SpDense SpDense k = k SpDense (Inj id) (Inj id) (\a b -> EPlus ext t a b)

-- handle absents
sparsePlusS SF _ _ SpAbsent sp2 k = k sp2 Noinj (Inj id) (\_ b -> b)
sparsePlusS ST _ t SpAbsent sp2 k =
  k (SpSparse sp2) (Inj $ \_ -> ENothing ext (applySparse sp2 (fromSMTy t))) (Inj $ EJust ext) (\_ b -> EJust ext b)

sparsePlusS _ SF _ sp1 SpAbsent k = k sp1 (Inj id) Noinj (\a _ -> a)
sparsePlusS _ ST t sp1 SpAbsent k =
  k (SpSparse sp1) (Inj $ EJust ext) (Inj $ \_ -> ENothing ext (applySparse sp1 (fromSMTy t))) (\a _ -> EJust ext a)

-- double sparse yields sparse
sparsePlusS _ _ t (SpSparse sp1) (SpSparse sp2) k =
  sparsePlusS ST ST t sp1 sp2 $ \sp3 (Inj inj1) (Inj inj2) plus ->
    k (SpSparse sp3)
      (Inj $ \a -> emaybe a (ENothing ext (applySparse sp3 (fromSMTy t))) (EJust ext (inj1 (evar IZ))))
      (Inj $ \b -> emaybe b (ENothing ext (applySparse sp3 (fromSMTy t))) (EJust ext (inj2 (evar IZ))))
      (\a b ->
        elet b $
          emaybe (weakenExpr WSink a)
            (emaybe (evar IZ)
              (ENothing ext (applySparse sp3 (fromSMTy t)))
              (EJust ext (inj2 (evar IZ))))
            (emaybe (evar (IS IZ))
              (EJust ext (inj1 (evar IZ)))
              (EJust ext (plus (evar (IS IZ)) (evar IZ)))))

-- single sparse can yield non-sparse if the other argument is always present
sparsePlusS SF _ t (SpSparse sp1) sp2 k =
  sparsePlusS SF ST t sp1 sp2 $ \sp3 _ (Inj inj2) plus ->
    k sp3 Noinj (Inj inj2)
      (\a b ->
        elet b $
          emaybe (weakenExpr WSink a)
            (inj2 (evar IZ))
            (plus (evar IZ) (evar (IS IZ))))
sparsePlusS ST _ t (SpSparse sp1) sp2 k =
  sparsePlusS ST ST t sp1 sp2 $ \sp3 (Inj inj1) (Inj inj2) plus ->
    k (SpSparse sp3)
      (Inj $ \a -> emaybe a (ENothing ext (applySparse sp3 (fromSMTy t))) (EJust ext (inj1 (evar IZ))))
      (Inj $ \b -> EJust ext (inj2 b))
      (\a b ->
        elet b $
          emaybe (weakenExpr WSink a)
            (EJust ext (inj2 (evar IZ)))
            (EJust ext (plus (evar IZ) (evar (IS IZ)))))
sparsePlusS req1 req2 t sp1 (SpSparse sp2) k =
  sparsePlusS req2 req1 t (SpSparse sp2) sp1 $ \sp3 inj1 inj2 plus ->
    k sp3 inj2 inj1 (flip plus)

-- products
sparsePlusS req1 req2 (SMTPair ta tb) (SpPair sp1a sp1b) (SpPair sp2a sp2b) k =
  sparsePlusS req1 req2 ta sp1a sp2a $ \sp3a minj13a minj23a plusa ->
  sparsePlusS req1 req2 tb sp1b sp2b $ \sp3b minj13b minj23b plusb ->
    k (SpPair sp3a sp3b)
      (withInj2 minj13a minj13b $ \inj13a inj13b ->
        \x1 -> eunPair x1 $ \_ x1a x1b -> EPair ext (inj13a x1a) (inj13b x1b))
      (withInj2 minj23a minj23b $ \inj23a inj23b ->
        \x2 -> eunPair x2 $ \_ x2a x2b -> EPair ext (inj23a x2a) (inj23b x2b))
      (\x1 x2 ->
        eunPair x1 $ \w1 x1a x1b ->
        eunPair (weakenExpr w1 x2) $ \w2 x2a x2b ->
          EPair ext (plusa (weakenExpr w2 x1a) x2a) (plusb (weakenExpr w2 x1b) x2b))
sparsePlusS req1 req2 t sp1@SpPair{} SpDense k = sparsePlusS req1 req2 t sp1 (SpPair SpDense SpDense) k
sparsePlusS req1 req2 t SpDense sp2@SpPair{} k = sparsePlusS req1 req2 t (SpPair SpDense SpDense) sp2 k

-- coproducts
sparsePlusS _ _ (SMTLEither ta tb) (SpLEither sp1a sp1b) (SpLEither sp2a sp2b) k =
  sparsePlusS ST ST ta sp1a sp2a $ \(sp3a :: Sparse _t3 t3a) (Inj inj13a) (Inj inj23a) plusa ->
  sparsePlusS ST ST tb sp1b sp2b $ \(sp3b :: Sparse _t3' t3b) (Inj inj13b) (Inj inj23b) plusb ->
    let nil :: Ex e (TLEither t3a t3b) ; nil = ELNil ext (applySparse sp3a (fromSMTy ta)) (applySparse sp3b (fromSMTy tb))
        inl :: Ex e t3a -> Ex e (TLEither t3a t3b) ; inl = ELInl ext (applySparse sp3b (fromSMTy tb))
        inr :: Ex e t3b -> Ex e (TLEither t3a t3b) ; inr = ELInr ext (applySparse sp3a (fromSMTy ta))
    in
    k (SpLEither sp3a sp3b)
      (Inj $ \x1 -> elcase x1 nil (inl (inj13a (evar IZ))) (inr (inj13b (evar IZ))))
      (Inj $ \x2 -> elcase x2 nil (inl (inj23a (evar IZ))) (inr (inj23b (evar IZ))))
      (\x1 x2 ->
        elet x2 $
          elcase (weakenExpr WSink x1)
            (elcase (evar IZ)
              nil
              (inl (inj23a (evar IZ)))
              (inr (inj23b (evar IZ))))
            (elcase (evar (IS IZ))
              (inl (inj13a (evar IZ)))
              (inl (plusa (evar (IS IZ)) (evar IZ)))
              (EError ext (applySparse (SpLEither sp3a sp3b) (fromSMTy (SMTLEither ta tb))) "plusS ll+lr"))
            (elcase (evar (IS IZ))
              (inr (inj13b (evar IZ)))
              (EError ext (applySparse (SpLEither sp3a sp3b) (fromSMTy (SMTLEither ta tb))) "plusS lr+ll")
              (inr (plusb (evar (IS IZ)) (evar IZ)))))
sparsePlusS req1 req2 t sp1@SpLEither{} SpDense k = sparsePlusS req1 req2 t sp1 (SpLEither SpDense SpDense) k
sparsePlusS req1 req2 t SpDense sp2@SpLEither{} k = sparsePlusS req1 req2 t (SpLEither SpDense SpDense) sp2 k

-- coproducts with partially known arguments: if we have a non-nil
-- always-present coproduct argument, the result is dense, otherwise we
-- introduce sparsity
sparsePlusS _ SF (SMTLEither ta _) (SpLeft sp1a) (SpLEither sp2a _) k =
  sparsePlusS ST SF ta sp1a sp2a $ \sp3a (Inj inj13a) _ plusa ->
    k (SpLeft sp3a)
      (Inj inj13a)
      Noinj
      (\x1 x2 ->
        elet x1 $
          elcase (weakenExpr WSink x2)
            (inj13a (evar IZ))
            (plusa (evar (IS IZ)) (evar IZ))
            (EError ext (applySparse sp3a (fromSMTy ta)) "plusS !ll+lr"))

sparsePlusS _ ST (SMTLEither ta _) (SpLeft sp1a) (SpLEither sp2a _) k =
  sparsePlusS ST ST ta sp1a sp2a $ \sp3a (Inj inj13a) (Inj inj23a) plusa ->
    k (SpSparse (SpLeft sp3a))
      (Inj $ \x1 -> EJust ext (inj13a x1))
      (Inj $ \x2 ->
        elcase x2
          (ENothing ext (applySparse sp3a (fromSMTy ta)))
          (EJust ext (inj23a (evar IZ)))
          (EError ext (STMaybe (applySparse sp3a (fromSMTy ta))) "plusSi2 !ll+lr"))
      (\x1 x2 ->
        elet x1 $
          EJust ext $
            elcase (weakenExpr WSink x2)
              (inj13a (evar IZ))
              (plusa (evar (IS IZ)) (evar IZ))
              (EError ext (applySparse sp3a (fromSMTy ta)) "plusS !ll+lr"))

sparsePlusS req1 req2 t sp1@SpLEither{} sp2@SpLeft{} k =
  sparsePlusS req2 req1 t sp2 sp1 $ \sp3a inj13a inj23a plusa -> k sp3a inj23a inj13a (flip plusa)
sparsePlusS req1 req2 t sp1@SpLeft{} SpDense k = sparsePlusS req1 req2 t sp1 (SpLEither SpDense SpDense) k
sparsePlusS req1 req2 t SpDense sp2@SpLeft{} k = sparsePlusS req1 req2 t (SpLEither SpDense SpDense) sp2 k

sparsePlusS _ SF (SMTLEither _ tb) (SpRight sp1b) (SpLEither _ sp2b) k =
  sparsePlusS ST SF tb sp1b sp2b $ \sp3b (Inj inj13b) _ plusb ->
    k (SpRight sp3b)
      (Inj inj13b)
      Noinj
      (\x1 x2 ->
        elet x1 $
          elcase (weakenExpr WSink x2)
            (inj13b (evar IZ))
            (EError ext (applySparse sp3b (fromSMTy tb)) "plusS !lr+ll")
            (plusb (evar (IS IZ)) (evar IZ)))

sparsePlusS _ ST (SMTLEither _ tb) (SpRight sp1b) (SpLEither _ sp2b) k =
  sparsePlusS ST ST tb sp1b sp2b $ \sp3b (Inj inj13b) (Inj inj23b) plusb ->
    k (SpSparse (SpRight sp3b))
      (Inj $ \x1 -> EJust ext (inj13b x1))
      (Inj $ \x2 ->
        elcase x2
          (ENothing ext (applySparse sp3b (fromSMTy tb)))
          (EError ext (STMaybe (applySparse sp3b (fromSMTy tb))) "plusSi2 !lr+ll")
          (EJust ext (inj23b (evar IZ))))
      (\x1 x2 ->
        elet x1 $
          EJust ext $
            elcase (weakenExpr WSink x2)
              (inj13b (evar IZ))
              (EError ext (applySparse sp3b (fromSMTy tb)) "plusS !lr+ll")
              (plusb (evar (IS IZ)) (evar IZ)))

sparsePlusS req1 req2 t sp1@SpLEither{} sp2@SpRight{} k =
  sparsePlusS req2 req1 t sp2 sp1 $ \sp3b inj13b inj23b plusb -> k sp3b inj23b inj13b (flip plusb)
sparsePlusS req1 req2 t sp1@SpRight{} SpDense k = sparsePlusS req1 req2 t sp1 (SpLEither SpDense SpDense) k
sparsePlusS req1 req2 t SpDense sp2@SpRight{} k = sparsePlusS req1 req2 t (SpLEither SpDense SpDense) sp2 k

-- dense same-branch coproducts simply recurse
sparsePlusS req1 req2 (SMTLEither ta _) (SpLeft sp1) (SpLeft sp2) k =
  sparsePlusS req1 req2 ta sp1 sp2 $ \sp3 inj1 inj2 plus ->
    k (SpLeft sp3) inj1 inj2 plus
sparsePlusS req1 req2 (SMTLEither _ tb) (SpRight sp1) (SpRight sp2) k =
  sparsePlusS req1 req2 tb sp1 sp2 $ \sp3 inj1 inj2 plus ->
    k (SpRight sp3) inj1 inj2 plus

-- dense, mismatched coproducts are valid as long as we don't actually invoke
-- plus at runtime (injections are fine)
sparsePlusS SF SF _ SpLeft{} SpRight{} k =
  k SpAbsent Noinj Noinj (\_ _ -> EError ext STNil "plusS !ll+!lr")
sparsePlusS SF ST (SMTLEither _ tb) SpLeft{} (SpRight sp2) k =
  k (SpRight sp2) Noinj (Inj id)
    (\_ _ -> EError ext (applySparse sp2 (fromSMTy tb)) "plusS !ll+?lr")
sparsePlusS ST SF (SMTLEither ta _) (SpLeft sp1) SpRight{} k =
  k (SpLeft sp1) (Inj id) Noinj
    (\_ _ -> EError ext (applySparse sp1 (fromSMTy ta)) "plusS !lr+?ll")
sparsePlusS ST ST (SMTLEither ta tb) (SpLeft sp1) (SpRight sp2) k =
  -- note: we know that this cannot be ELNil, but the returned 'Sparse' unfortunately claims to allow it.
  k (SpLEither sp1 sp2)
    (Inj $ \a -> ELInl ext (applySparse sp2 (fromSMTy tb)) a)
    (Inj $ \b -> ELInr ext (applySparse sp1 (fromSMTy ta)) b)
    (\_ _ -> EError ext (STLEither (applySparse sp1 (fromSMTy ta)) (applySparse sp2 (fromSMTy tb))) "plusS ?ll+?lr")

sparsePlusS req1 req2 t sp1@SpRight{} sp2@SpLeft{} k =  -- the errors are not flipped, but eh
  sparsePlusS req2 req1 t sp2 sp1 $ \sp3 inj1 inj2 plus -> k sp3 inj2 inj1 (flip plus)

-- maybe
sparsePlusS _ _ (SMTMaybe t) (SpMaybe sp1) (SpMaybe sp2) k =
  sparsePlusS ST ST t sp1 sp2 $ \sp3 (Inj inj1) (Inj inj2) plus ->
    k (SpMaybe sp3)
      (Inj $ \a -> emaybe a (ENothing ext (applySparse sp3 (fromSMTy t))) (EJust ext (inj1 (evar IZ))))
      (Inj $ \b -> emaybe b (ENothing ext (applySparse sp3 (fromSMTy t))) (EJust ext (inj2 (evar IZ))))
      (\a b ->
        elet b $
          emaybe (weakenExpr WSink a)
            (emaybe (evar IZ)
              (ENothing ext (applySparse sp3 (fromSMTy t)))
              (EJust ext (inj2 (evar IZ))))
            (emaybe (evar (IS IZ))
              (EJust ext (inj1 (evar IZ)))
              (EJust ext (plus (evar (IS IZ)) (evar IZ)))))
sparsePlusS req1 req2 t sp1@SpMaybe{} SpDense k = sparsePlusS req1 req2 t sp1 (SpMaybe SpDense) k
sparsePlusS req1 req2 t SpDense sp2@SpMaybe{} k = sparsePlusS req1 req2 t (SpMaybe SpDense) sp2 k

-- maybe with partially known arguments: if we have an always-present Just
-- argument, the result is dense, otherwise we introduce sparsity by weakening
-- to SpMaybe
sparsePlusS _ SF (SMTMaybe t) (SpJust sp1) (SpMaybe sp2) k =
  sparsePlusS ST SF t sp1 sp2 $ \sp3 (Inj inj1) _ plus ->
    k (SpJust sp3)
      (Inj inj1)
      Noinj
      (\a b ->
        elet a $
          emaybe (weakenExpr WSink b)
            (inj1 (evar IZ))
            (plus (evar (IS IZ)) (evar IZ)))
sparsePlusS _ ST (SMTMaybe t) (SpJust sp1) (SpMaybe sp2) k =
  sparsePlusS ST ST t sp1 sp2 $ \sp3 (Inj inj1) (Inj inj2) plus ->
    k (SpMaybe sp3)
      (Inj $ \a -> EJust ext (inj1 a))
      (Inj $ \b -> emaybe b (ENothing ext (applySparse sp3 (fromSMTy t))) (EJust ext (inj2 (evar IZ))))
      (\a b ->
        elet a $
          emaybe (weakenExpr WSink b)
            (EJust ext (inj1 (evar IZ)))
            (EJust ext (plus (evar (IS IZ)) (evar IZ))))

sparsePlusS req1 req2 t sp1@SpMaybe{} sp2@SpJust{} k =
  sparsePlusS req2 req1 t sp2 sp1 $ \sp3 inj2 inj1 plus -> k sp3 inj1 inj2 (flip plus)
sparsePlusS req1 req2 t sp1@SpJust{} SpDense k = sparsePlusS req1 req2 t sp1 (SpMaybe SpDense) k
sparsePlusS req1 req2 t SpDense sp2@SpJust{} k = sparsePlusS req1 req2 t (SpMaybe SpDense) sp2 k

-- dense same-branch maybes simply recurse
sparsePlusS req1 req2 (SMTMaybe t) (SpJust sp1) (SpJust sp2) k =
  sparsePlusS req1 req2 t sp1 sp2 $ \sp3 inj1 inj2 plus ->
    k (SpJust sp3) inj1 inj2 plus

-- dense array cotangents simply recurse
sparsePlusS req1 req2 (SMTArr _ t) (SpArr sp1) (SpArr sp2) k =
  sparsePlusS req1 req2 t sp1 sp2 $ \sp3 minj1 minj2 plus ->
    k (SpArr sp3)
      (withInj minj1 $ \inj1 -> emap (inj1 (EVar ext (applySparse sp1 (fromSMTy t)) IZ)))
      (withInj minj2 $ \inj2 -> emap (inj2 (EVar ext (applySparse sp2 (fromSMTy t)) IZ)))
      (ezipWith (plus (EVar ext (applySparse sp1 (fromSMTy t)) (IS IZ))
                      (EVar ext (applySparse sp2 (fromSMTy t)) IZ)))
sparsePlusS req1 req2 t (SpArr sp1) SpDense k = sparsePlusS req1 req2 t (SpArr sp1) (SpArr SpDense) k
sparsePlusS req1 req2 t SpDense (SpArr sp2) k = sparsePlusS req1 req2 t (SpArr SpDense) (SpArr sp2) k