1 files changed, 16 insertions, 10 deletions
diff --git a/src/Interpreter/Accum.hs b/src/Interpreter/Accum.hs
index b6a91df..af7be1e 100644
--- a/src/Interpreter/Accum.hs
+++ b/src/Interpreter/Accum.hs
@@ -51,6 +51,7 @@ type family Rep' s t where
   Rep' s TNil = ()
   Rep' s (TPair a b) = (Rep' s a, Rep' s b)
   Rep' s (TEither a b) = Either (Rep' s a) (Rep' s b)
+  Rep' s (TMaybe t) = Maybe (Rep' s t)
   Rep' s (TArr n t) = Array n (Rep' s t)
   Rep' s (TScal sty) = ScalRep sty
   Rep' s (TAccum t) = Accum s t
@@ -61,16 +62,13 @@ data Accum s t = Accum (STy t) (ForeignPtr ())
 tSize :: Proxy s -> STy t -> Rep' s t -> Int
 tSize p ty x = tSize' p ty (Just x)
 
--- | Passing Nothing as the value means "this is (inside) an array element".
-tSize' :: Proxy s -> STy t -> Maybe (Rep' s t) -> Int
-tSize' p typ val = case typ of
+tSize' :: Proxy s -> STy t -> Int
+tSize' p typ = case typ of
   STNil -> 0
-  STPair a b -> tSize' p a (fst <$> val) + tSize' p b (snd <$> val)
-  STEither a b ->
-    case val of
-      Nothing -> 1 + max (tSize' p a Nothing) (tSize' p b Nothing)
-      Just (Left x) -> 1 + tSize' p a (Just x)   -- '1 +' is for runtime sanity checking
-      Just (Right y) -> 1 + tSize' p b (Just y)  -- idem
+  STPair a b -> tSize' p a + tSize' p b
+  STEither a b -> 1 + max (tSize' p a) (tSize' p b)
+  -- Representation of Maybe t is the same as Either () t; the add operation is different, however.
+  STMaybe t -> tSize' p (STEither STNil t)
   STArr ndim t ->
     case val of
       Nothing -> error "Nested arrays not supported in this implementation"
@@ -99,7 +97,7 @@ accumWrite (Accum topty fptr) top_value = withForeignPtr fptr $ \(Ptr addr#) ->
         go inarr b (snd val) off1
       STEither a b -> do
         let !(I# off#) = off
-        case val of
+        off1 <- case val of
           Left x -> do
             let !(I8# tag#) = 0
             writeInt8# addr# off# tag#
@@ -108,6 +106,11 @@ accumWrite (Accum topty fptr) top_value = withForeignPtr fptr $ \(Ptr addr#) ->
             let !(I8# tag#) = 1
             writeInt8# addr# off# tag#
             go inarr b y (off + 1)
+        if inarr
+          then return (off + 1 + max (tSize' (Proxy @s) a Nothing) (tSize' (Proxy @s) b Nothing))
+          else return off1
+      -- Representation is the same, but add operation is different
+      STMaybe t -> go inarr (STEither STNil t) (maybe (Left ()) Right val) off
       STArr _ t
         | inarr -> error "Nested arrays not supported in this implementation"
         | otherwise -> do
@@ -158,6 +161,8 @@ accumRead (Accum topty fptr) = AcM $ withForeignPtr fptr $ \(Ptr addr#) ->
         if inarr
           then return (off + 1 + max (tSize' (Proxy @s) a Nothing) (tSize' (Proxy @s) b Nothing), val)
           else return (off1, val)
+      -- Representation is the same, but add operation is different
+      STMaybe t -> second (either (const Nothing) Just) <$> go inarr (STEither STNil t)  off
       STArr ndim t
         | inarr -> error "Nested arrays not supported in this implementation"
         | otherwise -> do
@@ -219,6 +224,7 @@ accumAdd (Accum topty fptr) top_depth top_index top_value = AcM $ withForeignPtr
       (STEither t1 _, Left idx1, Left val1) -> go inarr t1 dep idx1 val1 off
       (STEither _ t2, Right idx2, Right val2) -> go inarr t2 dep idx2 val2 off
       (STEither{}, _, _) -> error "Mismatching idx/val for Either in accumAdd"
+      (STMaybe t, _, _) -> _ idx val
       (STArr rank eltty, _, _)
         | inarr -> error "Nested arrays"
         | otherwise -> do