Cata tactic should generalize let and ensure unifiability (#1938)

* Let must be generalized when doing a letForEach

* Don't attempt to cata terms which don't unify with any args

* Add a test

Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

Author: isovector

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

@@ -33,6 +33,7 @@ import           GhcPlugins (isSymOcc, mkVarOccFS)
 import           OccName (occName)
 import           PatSyn
 import           Type hiding (Var)
+import           TysPrim (alphaTy)
 import           Wingman.CodeGen.Utils
 import           Wingman.GHC
 import           Wingman.Judgements
@@ -309,7 +310,8 @@ letForEach rename solve (unHypothesis -> hy) jdg = do
       let g = jGoal jdg
       terms <- fmap sequenceA $ for hy $ \hi -> do
         let name = rename $ hi_name hi
-        res <- tacticToRule jdg $ solve hi
+        let generalized_let_ty = CType alphaTy
+        res <- tacticToRule (withNewGoal generalized_let_ty jdg) $ solve hi
         pure $ fmap ((name,) . unLoc) res
       let hy' = fmap (g <$) $ syn_val terms
           matches = fmap (fmap (\(occ, expr) -> valBind (occNameToStr occ) expr)) terms

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

@@ -6,7 +6,7 @@ import           Control.Applicative (empty)
 import           Control.Lens ((<>~))
 import           Control.Monad.Error.Class
 import           Control.Monad.Reader
-import           Control.Monad.State.Class (gets, modify)
+import           Control.Monad.State.Class (gets, modify, MonadState)
 import           Control.Monad.State.Strict (StateT (..), execStateT)
 import           Control.Monad.Trans.Maybe
 import           Data.Coerce
@@ -217,6 +217,20 @@ unify goal inst = do
     Nothing -> throwError (UnificationError inst goal)
 
 
+------------------------------------------------------------------------------
+-- | Attempt to unify two types.
+canUnify
+    :: MonadState TacticState m
+    => CType -- ^ The goal type
+    -> CType -- ^ The type we are trying unify the goal type with
+    -> m Bool
+canUnify goal inst = do
+  skolems <- gets ts_skolems
+  case tryUnifyUnivarsButNotSkolems skolems goal inst of
+    Just _ -> pure True
+    Nothing -> pure False
+
+
 ------------------------------------------------------------------------------
 -- | Prefer the first tactic to the second, if the bool is true. Otherwise, just run the second tactic.
 --
@@ -312,6 +326,17 @@ lookupNameInContext name = do
     Nothing      -> empty
 
 
+getDefiningType
+    :: (MonadError TacticError m, MonadReader Context m)
+    => m CType
+getDefiningType = do
+  calling_fun_name <- fst . head <$> asks ctxDefiningFuncs
+  maybe
+    (throwError $ NotInScope calling_fun_name)
+    pure
+      =<< lookupNameInContext calling_fun_name
+
+
 ------------------------------------------------------------------------------
 -- | Build a 'HyInfo' for an imported term.
 createImportedHyInfo :: OccName -> CType -> HyInfo CType

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

@@ -8,8 +8,10 @@ module Wingman.Tactics
 import           ConLike (ConLike(RealDataCon))
 import           Control.Applicative (Alternative(empty))
 import           Control.Lens ((&), (%~), (<>~))
+import           Control.Monad (filterM)
 import           Control.Monad (unless)
 import           Control.Monad.Except (throwError)
+import           Control.Monad.Extra (anyM)
 import           Control.Monad.Reader.Class (MonadReader (ask))
 import           Control.Monad.State.Strict (StateT(..), runStateT, gets)
 import           Data.Bool (bool)
@@ -475,21 +477,38 @@ nary n =
     mkInvForAllTys [alphaTyVar, betaTyVar] $
       mkFunTys' (replicate n alphaTy) betaTy
 
+
 self :: TacticsM ()
 self =
   fmap listToMaybe getCurrentDefinitions >>= \case
     Just (self, _) -> useNameFromContext apply self
     Nothing -> throwError $ TacticPanic "no defining function"
 
 
+------------------------------------------------------------------------------
+-- | Perform a catamorphism when destructing the given 'HyInfo'. This will
+-- result in let binding, making values that call the defining function on each
+-- destructed value.
 cata :: HyInfo CType -> TacticsM ()
 cata hi = do
+  (_, _, calling_args, _)
+      <- tacticsSplitFunTy . unCType <$> getDefiningType
+  freshened_args <- traverse freshTyvars calling_args
   diff <- hyDiff $ destruct hi
+
+  -- For for every destructed term, check to see if it can unify with any of
+  -- the arguments to the calling function. If it doesn't, we don't try to
+  -- perform a cata on it.
+  unifiable_diff <- flip filterM (unHypothesis diff) $ \hi ->
+    flip anyM freshened_args $ \ty ->
+      canUnify (hi_type hi) $ CType ty
+
   rule $
     letForEach
       (mkVarOcc . flip mappend "_c" . occNameString)
-      (\hi -> self >> commit (apply hi) assumption)
-      diff
+      (\hi -> self >> commit (assume $ hi_name hi) assumption)
+      $ Hypothesis unifiable_diff
+
 
 collapse :: TacticsM ()
 collapse = do

plugins/hls-tactics-plugin/test/CodeAction/RunMetaprogramSpec.hs

@@ -33,5 +33,6 @@ spec = do
     metaTest 11 11 "MetaUseMethod"
     metaTest  9 38 "MetaCataCollapse"
     metaTest  7 16 "MetaCataCollapseUnary"
+    metaTest 21 32 "MetaCataAST"
     metaTest  6 46 "MetaPointwise"
     metaTest  4 28 "MetaUseSymbol"

plugins/hls-tactics-plugin/test/golden/MetaCataAST.expected.hs

@@ -0,0 +1,34 @@
+{-# LANGUAGE GADTs #-}
+
+data AST a where
+    BoolLit :: Bool -> AST Bool
+    IntLit :: Int -> AST Int
+    If :: AST Bool -> AST a -> AST a -> AST a
+    Equal :: AST a -> AST a -> AST Bool
+
+eval :: AST a -> a
+-- NOTE(sandy): There is an unrelated bug that is shown off in this test
+-- namely, that
+--
+-- @eval (IntLit n) = _@
+--
+-- but should be
+--
+-- @eval (IntLit n) = n@
+--
+-- https://github.com/haskell/haskell-language-server/issues/1937
+
+eval (BoolLit b) = b
+eval (IntLit n) = _
+eval (If ast ast' ast_a)
+  = let
+      ast_c = eval ast
+      ast'_c = eval ast'
+      ast_a_c = eval ast_a
+    in _ ast_c ast'_c ast_a_c
+eval (Equal ast ast')
+  = let
+      ast_c = eval ast
+      ast'_c = eval ast'
+    in _ ast_c ast'_c
+

plugins/hls-tactics-plugin/test/golden/MetaCataAST.hs

@@ -0,0 +1,22 @@
+{-# LANGUAGE GADTs #-}
+
+data AST a where
+    BoolLit :: Bool -> AST Bool
+    IntLit :: Int -> AST Int
+    If :: AST Bool -> AST a -> AST a -> AST a
+    Equal :: AST a -> AST a -> AST Bool
+
+eval :: AST a -> a
+-- NOTE(sandy): There is an unrelated bug that is shown off in this test
+-- namely, that
+--
+-- @eval (IntLit n) = _@
+--
+-- but should be
+--
+-- @eval (IntLit n) = n@
+--
+-- https://github.com/haskell/haskell-language-server/issues/1937
+
+eval = [wingman| intros x, cata x; collapse |]
+