Add some missing cases

2 files changed, 18 insertions, 0 deletions

diff --git a/src/AST/Count.hs b/src/AST/Count.hs
index 40a46f6..ad68685 100644
--- a/src/AST/Count.hs
+++ b/src/AST/Count.hs
@@ -110,6 +110,9 @@ occCountGeneral onehot unpush alter many = go WId
       EInl _ _ e -> re e
       EInr _ _ e -> re e
       ECase _ e a b -> re e <> (re1 a `alter` re1 b)
+      ENothing _ _ -> mempty
+      EJust _ e -> re e
+      EMaybe _ a b e -> re a <> re1 b <> re e
       EConstArr{} -> mempty
       EBuild1 _ a b -> re a <> many (re1 b)
       EBuild _ _ a b -> re a <> many (re1 b)
diff --git a/src/AST/Pretty.hs b/src/AST/Pretty.hs
index f5e681a..7f60db1 100644
--- a/src/AST/Pretty.hs
+++ b/src/AST/Pretty.hs
@@ -92,6 +92,21 @@ ppExpr' d val = \case
       showString "case " . e' . showString (" of { Inl " ++ name1 ++ " -> ") . a'
       . showString (" ; Inr " ++ name2 ++ " -> ") . b' . showString " }"
 
+  ENothing _ _ -> return $ showString "nothing"
+
+  EJust _ e -> do
+    e' <- ppExpr' 11 val e
+    return $ showParen (d > 10) $ showString "Just " . e'
+
+  EMaybe _ a b e -> do
+    let STMaybe t = typeOf e
+    a' <- ppExpr' 11 val a
+    name <- genNameIfUsedIn t IZ b
+    b' <- ppExpr' 11 (Const name `SCons` val) b
+    e' <- ppExpr' 11 val e
+    return $ showParen (d > 10) $
+      showString "maybe " . a' . showString " " . b' . showString " " . e'
+
   EConstArr _ _ ty v
     | Dict <- scalRepIsShow ty -> return $ showsPrec d v