hack together some union method for sorts

src/Language/Fixpoint/SortCheck.hs

@@ -1187,7 +1187,7 @@ checkURel e s1 s2 = unless (b1 == b2) (throwErrorAt $ errRel e s1 s2)
     b1            = s1 == boolSort
     b2            = s2 == boolSort
 
---------------------------------------------------------------------------------
+-------------------------------------------------------------------
 -- | Sort Unification on Expressions
 --------------------------------------------------------------------------------
 unifyUF :: Env -> UF -> Maybe Expr -> Sort -> Sort -> CheckM UF
@@ -1449,8 +1449,8 @@ unifyVar f e θ !i !t
       Nothing        -> return (updateVar i t θ)
 
 unifyVarUF :: Env -> Maybe Expr -> UF -> Int -> Sort -> CheckM UF
-unifyVarUF _ _ uf !_ t@(FVar !j)
-  = return (Union.union uf j t)
+unifyVarUF _ _ uf !i (FVar !j)
+  = return (Union.union uf j (FVar i))
 
 unifyVarUF _ _ uf !i !t
   = return (Union.union uf i t)

src/Language/Fixpoint/Union.hs

@@ -2,36 +2,79 @@ module Language.Fixpoint.Union where
 import Data.HashMap.Strict (lookup, insert, HashMap, empty)
 import Prelude hiding (lookup)
 import Language.Fixpoint.Types.Sorts (Sort(..))
-next :: Sort -> Maybe Int 
-next (FVar i) = Just i
-next _ = Nothing 
-unionVals :: UF -> Sort -> Sort -> UF
-unionVals _ _ _ = error "todo"
--- unionVals ufM i s1 s2 = error "todo"
--- unionVals ufM _ FInt FInt = ufM
--- unionVals ufM _ FReal FReal = ufM
--- unionVals ufM _ FInt FReal = ufM
--- unionVals ufM _ FReal FInt = ufM
+
+-- unionVals :: UF Sort -> Int -> Sort -> Sort -> UF Sort
+--     unionVals ufM _ SInt SInt = ufM
+--     unionVals ufM _ SFloat SFloat = ufM
+--     unionVals ufM _ (SFVar j) s = Union.union ufM j s
+--     unionVals ufM _ s (SFVar j) = Union.union ufM j s
+--     unionVals u i (SFunc s1 s2) (SFunc s1' s2') = 
+--         let u' = unionFuncArgs u i s1 s1' in 
+--             unionFuncArgs u' i s2 s2'
+--     unionVals _ _ s1 s2 = error ("Cannot unify " ++ show s1 ++ " " ++ show s2)
+--     next s = case s of
+--         SFVar i -> Just i
+--         _ -> Nothing
+
+unionMany :: UF -> Int -> Sort -> Sort -> UF
+unionMany uf i s1 s2 = case (s1, s2) of 
+    (FVar i1, _) -> union uf i1 s2
+    (_, FVar i2) -> union uf i2 s1
+    (_, _) -> unionVals uf i s1 s2
+
+--------------------------------------------------------------------------------
+-- | union for sorts in union find
+--------------------------------------------------------------------------------
+unionVals :: UF -> Int -> Sort -> Sort -> UF
+--------------------------------------------------------------------------------
+unionVals uf _ s1 s2 
+  | isNumericSort s1 && isNumericSort s2 = uf
+  where
+    isNumericSort FReal = True
+    isNumericSort FNum  = True
+    isNumericSort FFrac = True
+    isNumericSort _     = False
+
+unionVals uf _ (FObj _) (FObj _) = uf
+unionVals uf _ (FVar i) s = union uf i s
+unionVals uf _ s (FVar i) = union uf i s
+unionVals uf i (FFunc s1 s2) (FFunc s1' s2') = 
+    let uf' = unionMany uf i s1 s1' in 
+        unionMany uf' i s2 s2'
+unionVals uf i (FApp s1 s2) (FApp s1' s2') = 
+    let uf' = unionMany uf i s1 s1' in 
+        unionMany uf' i s2 s2'
+unionVals uf i (FAbs _ s) (FAbs _ s') = unionMany uf i s s'
+unionVals uf _ (FTC _) (FTC _) = uf
+unionVals _ _ _ _ = error "Cannot unify"
+
 
 newtype UF = MkUF (HashMap Int Sort) deriving (Show)
 new :: UF
 new = MkUF empty
 union :: UF -> Int -> Sort -> UF
 union u@(MkUF ufM) tyv s =
     -- find the root for tyv 
-    let tyv_root = find (MkUF ufM) tyv in
+    let tyv_root = findWithIndex (MkUF ufM) tyv in
     case tyv_root of
             -- if tyv not in union find, insert
             Nothing -> MkUF (insert tyv s ufM)
             -- otherwise, unify the current sort with 
             -- the new one and insert that
-            Just s' -> unionVals u s s'
+            Just (i, s') -> unionVals u i s s'
+
+findWithIndex :: UF -> Int -> Maybe (Int, Sort)
+findWithIndex u@(MkUF ufM) k = do 
+    s <- lookup k ufM 
+    case s of 
+        FVar i -> findWithIndex u i
+        s' -> Just (k, s')
 
 find  :: UF -> Int -> Maybe Sort
 find (MkUF ufM) = f 
     where 
     f k = do 
         s <- lookup k ufM 
-        case next s of 
-            Nothing -> Just s
-            Just i -> f i
+        case s of 
+            FVar i -> f i
+            s' -> Just s'