Use type variable names when comparing types

Refs #293

src/SearchIndex.hs

@@ -17,9 +17,11 @@ module SearchIndex
 
 import Import.NoFoundation
 import Control.Parallel.Strategies (Strategy, evalTraversable, rdeepseq)
+import Control.Monad.Trans.Writer (WriterT(..), tell)
 import Data.Trie (Trie)
 import Data.Version (Version)
 import qualified Data.Map as Map
+import qualified Data.Set as Set
 import qualified Data.Text as T
 import qualified Data.Trie as Trie
 import qualified Text.Parsec.Combinator as Parsec
@@ -340,44 +342,73 @@ tryStripPrefix pre s = fromMaybe s (T.stripPrefix pre s)
 -- >>> compare "Int" "a"
 -- Just (Just 1)
 --
+-- (The idea here being it's ok to show a more general version of the query,
+-- but usually not helpful to show a more concrete version of it.)
+--
 compareTypes :: P.Type -> P.Type -> Maybe Int
-compareTypes (P.TypeVar _) (P.TypeVar _) = Just 0
-compareTypes t (P.TypeVar _) = Just (1 + typeComplexity t)
-compareTypes (P.TypeLevelString s1) (P.TypeLevelString s2) | s1 == s2 = Just 0
-compareTypes (P.TypeWildcard _) t = Just (typeComplexity t)
-compareTypes (P.TypeConstructor q1) (P.TypeConstructor q2) | compareQual q1 q2 = Just 0
--- There is a special case for functions, since if the user _asked_ for a function,
--- they probably don't want to see something more general of type 'f a' or 'f a b'.
-compareTypes (P.TypeApp a b) (P.TypeApp c d)
-  | not (isFunction a) || isFunction c = (+) <$> compareTypes a c <*> compareTypes b d
-compareTypes (P.ForAll _ t1 _) t2 = compareTypes t1 t2
-compareTypes t1 (P.ForAll _ t2 _) = compareTypes t1 t2
-compareTypes (P.ConstrainedType _ t1) t2 = compareTypes t1 t2
-compareTypes t1 (P.ConstrainedType _ t2) = compareTypes t1 t2
-compareTypes P.REmpty P.REmpty = Just 0
-compareTypes t1@P.RCons{} t2 = compareRows t1 t2
-compareTypes t1 t2@P.RCons{} = compareRows t1 t2
-compareTypes (P.KindedType t1 _) t2 = compareTypes t1 t2
-compareTypes t1 (P.KindedType t2 _) = compareTypes t1 t2
--- Really, we should desugar any type operators here.
--- Since type operators are not supported in search right now, this is fine,
--- since we only care about functions, which are already in the correct
--- order as they come out of the parser.
-compareTypes (P.ParensInType t1) t2 = compareTypes t1 t2
-compareTypes t1 (P.ParensInType t2) = compareTypes t1 t2
-compareTypes _ _ = Nothing
+compareTypes type1 type2 =
+  map calculate . runWriterT $ go type1 type2
+  where
+  calculate :: (Int, [(Text, Text)]) -> Int
+  calculate (score, vars) = (10 * score) + typeVarPenalty vars
+
+  go :: P.Type -> P.Type -> WriterT [(Text, Text)] Maybe Int
+  go (P.TypeVar v1) (P.TypeVar v2) = tell [(v1, v2)] *> pure 0
+  go t (P.TypeVar _) = pure (1 + typeComplexity t)
+  go (P.TypeLevelString s1) (P.TypeLevelString s2) | s1 == s2 = pure 0
+  go (P.TypeWildcard _) t = pure (typeComplexity t)
+  go (P.TypeConstructor q1) (P.TypeConstructor q2) | compareQual q1 q2 = pure 0
+-- There is a special case for functions, since if the user _asked_ for a
+-- function, they probably don't want to see something more general of type 'f
+-- a' or 'f a b'.
+  go (P.TypeApp a b) (P.TypeApp c d)
+    | not (isFunction a) || isFunction c = (+) <$> go a c <*> go b d
+  go (P.ForAll _ t1 _) t2 = go t1 t2
+  go t1 (P.ForAll _ t2 _) = go t1 t2
+  go (P.ConstrainedType _ t1) t2 = go t1 t2
+  go t1 (P.ConstrainedType _ t2) = go t1 t2
+  go P.REmpty P.REmpty = pure 0
+  go t1@P.RCons{} t2 = goRows t1 t2
+  go t1 t2@P.RCons{} = goRows t1 t2
+  go (P.KindedType t1 _) t2 = go t1 t2
+  go t1 (P.KindedType t2 _) = go t1 t2
+  -- Really, we should desugar any type operators here.
+  -- Since type operators are not supported in search right now, this is fine,
+  -- since we only care about functions, which are already in the correct
+  -- order as they come out of the parser.
+  go (P.ParensInType t1) t2 = go t1 t2
+  go t1 (P.ParensInType t2) = go t1 t2
+  go _ _ = lift Nothing
+
+  goRows :: P.Type -> P.Type -> WriterT [(Text, Text)] Maybe Int
+  goRows r1 r2 = sum <$>
+    sequence [ go t1 t2
+             | (name, t1) <- fst (P.rowToList r1)
+             , (name', t2) <- fst (P.rowToList r2)
+             , name == name'
+             ]
+
+  -- Calculate a penalty based on the extent to which the type variables match.
+  -- Where differences occur, those which make the result more general than the
+  -- query are not penalised as harshly as those which make the result less
+  -- general than the query.
+  typeVarPenalty :: [(Text, Text)] -> Int
+  typeVarPenalty list =
+    penalty list + (3 * penalty (map swap list))
+    where
+    penalty =
+      map (second Set.singleton)
+      >>> Map.fromListWith Set.union
+      >>> Map.elems
+      -- If one element of the fsts is paired with more than one element of the
+      -- snds, penalise based on how many more elements of the snds there are.
+      >>> map (\s -> Set.size s - 1)
+      >>> sum
 
 isFunction :: P.Type -> Bool
 isFunction (P.TypeConstructor (P.Qualified _ (P.ProperName "Function"))) = True
 isFunction _ = False
 
-compareRows :: P.Type -> P.Type -> Maybe Int
-compareRows r1 r2 = sum <$>
-  sequence [ compareTypes t1 t2
-           | (name, t1) <- fst (P.rowToList r1)
-           , (name', t2) <- fst (P.rowToList r2)
-           , name == name'
-           ]
 
 typeComplexity :: P.Type -> Int
 typeComplexity (P.TypeApp a b) = 1 + typeComplexity a + typeComplexity b

test/SearchSpec.hs

@@ -1,11 +1,13 @@
 module SearchSpec where
 
 import Import.NoFoundation
+import Data.Maybe (fromJust)
 import Test.Hspec
 import Test.Hspec.QuickCheck (prop)
 import Test.QuickCheck as QC
 
 import Handler.Search (interleave)
+import SearchIndex (compareTypes, parseType)
 
 data Whatev = A | B | C | D | E
   deriving (Show, Eq, Enum, Bounded)
@@ -20,25 +22,102 @@ genSortedAssocList :: QC.Gen [(Whatev, Int)]
 genSortedAssocList =
   fmap (sortWith snd) genAssocList
 
+-- A version of 'expectationFailure' with its type relaxed from IO () to IO a.
+failed :: String -> IO a
+failed msg =
+  expectationFailure msg
+  >> error "should not be reached; hack to relax the type of expectationFailure"
+
+shouldBeLessThan :: (Show a, Ord a) => a -> a -> Expectation
+shouldBeLessThan x y =
+  if x < y
+    then pure ()
+    else expectationFailure
+           (show x <> " should have been less than " <> show y)
+
+shouldBeGreaterThan :: (Show a, Ord a) => a -> a -> Expectation
+shouldBeGreaterThan x y =
+  if x > y
+    then pure ()
+    else expectationFailure
+           (show x <> " should have been greater than " <> show y)
+
 spec :: Spec
 spec = do
-  prop "interleave is associative" $
-    forAll ((,,) <$> genAssocList <*> genAssocList <*> genAssocList) $
-      \(xs, ys, zs) ->
-        (xs `interleave` ys) `interleave` zs ===
-        xs `interleave` (ys `interleave` zs)
-
-  prop "[] is identity for interleave" $
-    forAll genAssocList $ \xs ->
-      QC.conjoin
-        [ xs `interleave` [] === xs
-        , [] `interleave` xs === xs
-        ]
-
-  prop "interleave preserves ordering" $
-    forAll ((,) <$> genSortedAssocList <*> genSortedAssocList) $
-      \(xs, ys) ->
-        let
-          zs = xs `interleave` ys
-        in
-          zs === sortWith snd zs
+  describe "interleave" $ do
+    prop "is associative" $
+      forAll ((,,) <$> genAssocList <*> genAssocList <*> genAssocList) $
+        \(xs, ys, zs) ->
+          (xs `interleave` ys) `interleave` zs ===
+          xs `interleave` (ys `interleave` zs)
+
+    prop "has identity []" $
+      forAll genAssocList $ \xs ->
+        QC.conjoin
+          [ xs `interleave` [] === xs
+          , [] `interleave` xs === xs
+          ]
+
+    prop "preserves ordering" $
+      forAll ((,) <$> genSortedAssocList <*> genSortedAssocList) $
+        \(xs, ys) ->
+          let
+            zs = xs `interleave` ys
+          in
+            zs === sortWith snd zs
+
+  describe "compareTypes" $ do
+    let
+      p = fromJust . parseType
+
+      -- Assert that two types should return a possible match with
+      -- compareTypes, and return the result.
+      shouldMatch t1 t2 =
+        maybe
+          (failed "should have returned a possible match")
+          pure
+          (compareTypes t1 t2)
+
+    it "respects alpha-equivalence" $ do
+      let
+        query = p "a -> b -> c -> b -> a"
+        cand0 = p "d -> e -> f -> e -> d"
+
+      shouldMatch query query >>= (`shouldBe` 0)
+      shouldMatch query cand0 >>= (`shouldBe` 0)
+
+    describe "with just type variables" $ do
+      let
+        query =           p "a -> b -> c -> b -> a"
+        moreGeneral =     p "d -> e -> f -> e -> h"
+        evenMoreGeneral = p "d -> e -> f -> g -> h"
+        lessGeneral =     p "a -> b -> b -> b -> a"
+        evenLessGeneral = p "a -> a -> a -> a -> a"
+
+      it "prioritises closer matches (more general)" $ do
+        x <- shouldMatch query moreGeneral
+        y <- shouldMatch query evenMoreGeneral
+        x `shouldBeLessThan` y
+
+      it "prioritises closer matches (less general)" $ do
+        x <- shouldMatch query lessGeneral
+        y <- shouldMatch query evenLessGeneral
+        x `shouldBeLessThan` y
+
+      it "prefers more general results to less general ones" $ do
+        x <- shouldMatch query evenMoreGeneral
+        y <- shouldMatch query lessGeneral
+        x `shouldBeLessThan` y
+
+    it "with both type variables and concrete types" $ do
+      let
+        query = p "a -> b -> Int"
+        cand0 = p "a -> a -> Int"
+        cand1 = p "a -> a -> b"
+
+      shouldMatch query query >>= (`shouldBe` 0)
+      shouldMatch query cand0 >>= (`shouldBeGreaterThan` 0)
+
+      x <- shouldMatch query cand0
+      y <- shouldMatch query cand1
+      x `shouldBeLessThan` y