Wingman: Add "New Unification Variable" helper (#2164)

* Add a newUnivar machinery, and strengthen unification tests

* Get rid of explicit univar constructions

* Cleanup AutoTupleSpec

Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Co-authored-by: Javier Neira <atreyu.bbb@gmail.com>

Author: 3 people

plugins/hls-tactics-plugin/src/Wingman/Machinery.hs

@@ -33,8 +33,9 @@ import           Refinery.Tactic.Internal
 import           System.Timeout (timeout)
 import           TcType
 import           Type (tyCoVarsOfTypeWellScoped)
+import           TysPrim (alphaTyVar, alphaTy)
 import           Wingman.Context (getInstance)
-import           Wingman.GHC (tryUnifyUnivarsButNotSkolems, updateSubst, tacticsGetDataCons)
+import           Wingman.GHC (tryUnifyUnivarsButNotSkolems, updateSubst, tacticsGetDataCons, freshTyvars)
 import           Wingman.Judgements
 import           Wingman.Simplify (simplify)
 import           Wingman.Types
@@ -229,6 +230,12 @@ newtype Reward a = Reward a
   deriving (Eq, Ord, Show) via a
 
 
+------------------------------------------------------------------------------
+-- | Generate a unique unification variable.
+newUnivar :: MonadState TacticState m => m Type
+newUnivar = do
+  freshTyvars $
+    mkInvForAllTys [alphaTyVar] alphaTy
 
 
 ------------------------------------------------------------------------------

plugins/hls-tactics-plugin/src/Wingman/Tactics.hs

@@ -36,7 +36,6 @@ import           Refinery.Tactic
 import           Refinery.Tactic.Internal
 import           TcType
 import           Type hiding (Var)
-import           TysPrim (betaTy, alphaTy, betaTyVar, alphaTyVar)
 import           Wingman.CodeGen
 import           Wingman.GHC
 import           Wingman.Judgements
@@ -543,10 +542,10 @@ applyByType ty = tracing ("applyByType " <> show ty) $ do
 -- | Make an n-ary function call of the form
 -- @(_ :: forall a b. a -> a -> b) _ _@.
 nary :: Int -> TacticsM ()
-nary n =
-  applyByType $
-    mkInvForAllTys [alphaTyVar, betaTyVar] $
-      mkFunTys' (replicate n alphaTy) betaTy
+nary n = do
+  a <- newUnivar
+  b <- newUnivar
+  applyByType $ mkFunTys' (replicate n a) b
 
 
 self :: TacticsM ()
@@ -589,8 +588,7 @@ letBind occs = do
           -> fmap (occ, )
            $ fmap (<$ jdg)
            $ fmap CType
-           $ freshTyvars
-           $ mkInvForAllTys [alphaTyVar] alphaTy
+           $ newUnivar
   rule $ nonrecLet occ_tys
 
 
@@ -630,7 +628,7 @@ collapse = do
 with_arg :: TacticsM ()
 with_arg = rule $ \jdg -> do
   let g = jGoal jdg
-  fresh_ty <- freshTyvars $ mkInvForAllTys [alphaTyVar] alphaTy
+  fresh_ty <- newUnivar
   a <- newSubgoal $ withNewGoal (CType fresh_ty) jdg
   f <- newSubgoal $ withNewGoal (coerce mkFunTys' [fresh_ty] g) jdg
   pure $ fmap noLoc $ (@@) <$> fmap unLoc f <*> fmap unLoc a

plugins/hls-tactics-plugin/test/AutoTupleSpec.hs

@@ -4,13 +4,13 @@
 
 module AutoTupleSpec where
 
+import Control.Monad (replicateM)
+import Control.Monad.State (evalState)
 import Data.Either (isRight)
 import OccName (mkVarOcc)
 import System.IO.Unsafe
 import Test.Hspec
 import Test.QuickCheck
-import Type (mkTyVarTy)
-import TysPrim (alphaTyVars)
 import TysWiredIn (mkBoxedTupleTy)
 import Wingman.Judgements (mkFirstJudgement)
 import Wingman.Machinery
@@ -24,7 +24,8 @@ spec = describe "auto for tuple" $ do
     property $ do
       -- Pick some number of variables
       n <- choose (1, 7)
-      let vars = fmap mkTyVarTy $ take n alphaTyVars
+      let vars = flip evalState defaultTacticState
+               $ replicateM n newUnivar
       -- Pick a random ordering
       in_vars  <- shuffle vars
       -- Randomly associate them into tuple types

plugins/hls-tactics-plugin/test/UnificationSpec.hs

@@ -4,6 +4,8 @@
 module UnificationSpec where
 
 import           Control.Arrow
+import           Control.Monad (replicateM, join)
+import           Control.Monad.State (evalState)
 import           Data.Bool (bool)
 import           Data.Functor ((<&>))
 import           Data.Maybe (mapMaybe)
@@ -13,10 +15,9 @@ import           Data.Tuple (swap)
 import           TcType (substTy, tcGetTyVar_maybe)
 import           Test.Hspec
 import           Test.QuickCheck
-import           Type (mkTyVarTy)
-import           TysPrim (alphaTyVars)
 import           TysWiredIn (mkBoxedTupleTy)
 import           Wingman.GHC
+import           Wingman.Machinery (newUnivar)
 import           Wingman.Types
 
 
@@ -25,8 +26,11 @@ spec = describe "unification" $ do
   it "should be able to unify univars with skolems on either side of the equality" $ do
     property $ do
       -- Pick some number of unification vars and skolem
-      n <- choose (1, 1)
-      let (skolems, take n -> univars) = splitAt n $ fmap mkTyVarTy alphaTyVars
+      n <- choose (1, 20)
+      let (skolems, take n -> univars)
+            = splitAt n
+            $ flip evalState defaultTacticState
+            $ replicateM (n * 2) newUnivar
       -- Randomly pair them
       skolem_uni_pairs <-
         for (zip skolems univars) randomSwap
@@ -42,13 +46,25 @@ spec = describe "unification" $ do
                  (CType lhs)
                  (CType rhs) of
             Just subst ->
-              -- For each pair, running the unification over the univar should
-              -- result in the skolem
-              conjoin $ zip univars skolems <&> \(uni, skolem) ->
-                let substd = substTy subst uni
-                 in counterexample (show substd) $
-                    counterexample (show skolem) $
-                      CType substd === CType skolem
+              conjoin $ join $
+                [ -- For each pair, running the unification over the univar should
+                  -- result in the skolem
+                  zip univars skolems <&> \(uni, skolem) ->
+                    let substd = substTy subst uni
+                    in counterexample (show substd) $
+                        counterexample (show skolem) $
+                          CType substd === CType skolem
+
+                  -- And also, no two univars should equal to one another
+                  -- before or after substitution.
+                , zip univars (tail univars) <&> \(uni1, uni2) ->
+                    let uni1_sub = substTy subst uni1
+                        uni2_sub = substTy subst uni2
+                    in counterexample (show uni1) $
+                        counterexample (show uni2) $
+                          CType uni1 =/= CType uni2 .&&.
+                          CType uni1_sub =/= CType uni2_sub
+                ]
             Nothing -> True === False