context elimination for variant types

haskell/src/SumTypes/EliminateContextArgument.hs

@@ -6,6 +6,9 @@ module SumTypes.EliminateContextArgument where
 
 import Prelude hiding (Monad(..), (>=), (<=), lookup)
 import GHC.Exts (Constraint)
+import qualified Data.Map as Map
+import Data.Maybe(isJust, fromJust)
+import Data.List(find)
 
 import SumTypes.Language
 import Util.ErrorMessages
@@ -42,6 +45,8 @@ instance WithTop () where
 instance (WithTop a, WithTop b) => WithTop (a, b) where
   top = (top, top)
 
+instance (WithTop v) => WithTop (Map.Map Name v) where
+  top = Map.empty
 
 -- Had to define an own monad type class.
 -- It is not possible otherwise to get the type constraint WithTop a.
@@ -85,6 +90,14 @@ matchSum :: Type -> String -> Infer (Type, Type)
 matchSum (Sum ty1 ty2) _ = return (ty1, ty2)
 matchSum ty err = fail [sumError ty err]
 
+matchVariant :: Type -> String -> Infer (Map.Map Name Type)
+matchVariant (Variant types) _ = return types
+matchVariant ty err = fail [variantError ty err]
+
+liftMaybe :: WithTop a => Maybe a -> String -> Infer a
+liftMaybe (Just a) _ = return a
+liftMaybe Nothing err = fail [err]
+
 -- try to inline lookup
 -- assume that ctx != Empty if not otherwise stated because this case was already covered in the previous step
 -- if there is a bind replace lookup definition with it
@@ -142,6 +155,19 @@ lookup t x = case parent t of
       (ty1, ty2) <- matchSum ety (show e)
       return ty2
     else lookup p x
+  Just p@(Tag n t' _) | t == t' -> do
+    lookup p x
+  Just p@(Match m cases _) | t == m -> do
+    lookup p x
+  Just p@(Match m cases _) | t /= m -> do
+    ety <- inferType m
+    typeMap <- matchVariant ety (show m)
+    let ml = find (\(_, x', t') -> t == t' && x == x') cases
+    if isJust ml
+    then do
+      let (l, _, _) = fromJust ml
+      liftMaybe (Map.lookup l typeMap) "Could not find labeled type"
+    else lookup p x
   Just p -> do
     lookup p x
   Nothing -> fail ["Unbound variable " ++ show x]
@@ -192,6 +218,22 @@ checkType p@(Case e n1 t1 n2 t2 _) = do
   (ty1, ty2) <- matchSum ety (show e)
   checkType t1
   checkType t2
+checkType p@(Tag n t _) = do
+  ty <- requiredType p
+  types <- matchVariant ty (show p)
+  let lty = Map.lookup n types
+  lty <- liftMaybe (Map.lookup n types) "Label not contained in Variant"
+  checkType t
+checkType p@(Match m cases _) = do
+  ety <- inferType m
+  types <- matchVariant ety (show m)
+  ty <- requiredType p
+  let subchecks =
+        map (\(l, x, t) -> do
+          lty <- liftMaybe (Map.lookup l types) "Could not find labeled type"
+          checkType t
+        ) cases
+  foldl (>>) (return ()) subchecks
 checkType t = do
   ty <- requiredType t
   ty' <- inferType t
@@ -231,6 +273,13 @@ requiredType t = case parent t of
   Just p@(Case e n1 t1 n2 t2 _) | t == t2 -> do
     ty <- requiredType p
     return ty
+  Just p@(Tag n t' _) | t == t' -> do
+    ty <- requiredType p
+    types <- matchVariant ty (show p)
+    liftMaybe (Map.lookup n types) "Label not contained in Variant"
+  Just p@(Match m cases _) -> do
+    ty <- requiredType p
+    return ty
   _ -> fail ["Wrong arg passed"]
 
 

haskell/test/SumTypes/SharedSpecs.hs

@@ -60,6 +60,8 @@ sharedSpec inferType = do
         let res = convert $ inferType tOkMatchWithDifferentOrder  in res `shouldBe` B.Inferred Nat
     it "should infer match where case shadows outer binding"$ do
         let res = convert $ inferType tOkMatchShadowBinding  in res `shouldBe` B.Inferred Nat
+    it "should infer match where bindings are and terms are the same" $ do
+        let res = convert $ inferType tOkCasesSameBinding  in res `shouldBe` B.Inferred (Fun Nat Nat)
 
     it "should fail while inferring an arithmetic expression with an unapplied lambda expression" $ do
         let res = isInferred $ convert $ inferType tFailArithmetic in res `shouldBe` False

haskell/test/SumTypes/TestCases.hs

@@ -29,6 +29,7 @@ tOkMatchSingle = root $ (anno Nat (match1 (anno (Variant (fromList [("a", Nat)])
 tOkMatchTwo = root $ (anno Nat (match2 (anno (Variant (fromList [("a", Nat), ("b", (Fun Nat Nat))])) (tag "b" (lam "x" (var "x")))) ("a", "x", succ $ var "x") ("b", "x", (app (var "x") zero))))
 tOkMatchWithDifferentOrder = root $ (anno Nat (match2 (anno (Variant (fromList [("b", (Fun Nat Nat)), ("a", Nat)])) (tag "b" (lam "x" (var "x")))) ("a", "x", succ $ var "x") ("b", "x", (app (var "x") zero))))
 tOkMatchShadowBinding = root $ (let' "x" zero (anno Nat (match2 (anno (Variant (fromList [("a", Nat), ("b", (Fun Nat Nat))])) (tag "b" (lam "x" (var "x")))) ("a", "x", succ $ var "x") ("b", "x", (app (var "x") zero)))))
+tOkCasesSameBinding = root $ (anno (Fun Nat Nat) (match2 (anno (Variant (fromList [("a", Nat), ("b", Nat)])) (tag "b" zero)) ("a", "x", (lam "y" (mult (var "y") (var "x")))) ("b", "x", (lam "y" (mult (var "y") (var "x"))))))
 
 
 tFailArithmetic = root $ (add (succ $ succ zero) (mult (anno (Fun Nat Nat) (lam "f" (var "f"))) (succ $ succ $ succ $ succ zero)))