1{-# LANGUAGE PatternGuards, ViewPatterns, FlexibleContexts, ScopedTypeVariables, TupleSections #-} 2{-# LANGUAGE GeneralizedNewtypeDeriving, DeriveFunctor #-} 3 4module GHC.Util.Unify( 5 Subst, fromSubst, 6 validSubst, removeParens, substitute, 7 unifyExp 8 ) where 9 10import Control.Applicative 11import Control.Monad 12import Data.Generics.Uniplate.DataOnly 13import Data.Char 14import Data.Data 15import Data.List.Extra 16import Util 17 18import GHC.Hs 19import GHC.Types.SrcLoc 20import GHC.Utils.Outputable hiding ((<>)) 21import GHC.Types.Name.Reader 22 23import Language.Haskell.GhclibParserEx.GHC.Hs.Pat 24import Language.Haskell.GhclibParserEx.GHC.Hs.Expr 25import Language.Haskell.GhclibParserEx.GHC.Utils.Outputable 26import Language.Haskell.GhclibParserEx.GHC.Types.Name.Reader 27import GHC.Util.HsExpr 28import GHC.Util.View 29import Data.Maybe 30import GHC.Data.FastString 31 32isUnifyVar :: String -> Bool 33isUnifyVar [x] = x == '?' || isAlpha x 34isUnifyVar [] = False 35isUnifyVar xs = all (== '?') xs 36 37--------------------------------------------------------------------- 38-- SUBSTITUTION DATA TYPE 39 40-- A list of substitutions. A key may be duplicated, you need to call 41-- 'check' to ensure the substitution is valid. 42newtype Subst a = Subst [(String, a)] 43 deriving (Semigroup, Monoid, Functor) 44 45-- Unpack the substitution. 46fromSubst :: Subst a -> [(String, a)] 47fromSubst (Subst xs) = xs 48 49instance Outputable a => Show (Subst a) where 50 show (Subst xs) = unlines [a ++ " = " ++ unsafePrettyPrint b | (a,b) <- xs] 51 52-- Check the unification is valid and simplify it. 53validSubst :: (a -> a -> Bool) -> Subst a -> Maybe (Subst a) 54validSubst eq = fmap Subst . mapM f . groupSort . fromSubst 55 where f (x, y : ys) | all (eq y) ys = Just (x, y) 56 f _ = Nothing 57 58-- Remove unnecessary brackets from a Subst. The first argument is a list of unification variables 59-- for which brackets should be removed from their substitutions. 60removeParens :: [String] -> Subst (LHsExpr GhcPs) -> Subst (LHsExpr GhcPs) 61removeParens noParens (Subst xs) = Subst $ 62 map (\(x, y) -> if x `elem` noParens then (x, fromParen y) else (x, y)) xs 63 64-- Peform a substition. 65-- Returns (suggested replacement, (refactor template, no bracket vars)). It adds/removes brackets 66-- for both the suggested replacement and the refactor template appropriately. The "no bracket vars" 67-- is a list of substituation variables which, when expanded, should have the brackets stripped. 68-- 69-- Examples: 70-- (traverse foo (bar baz), (traverse f (x), [])) 71-- (zipWith foo bar baz, (f a b, [f])) 72substitute :: Subst (LHsExpr GhcPs) -> LHsExpr GhcPs -> (LHsExpr GhcPs, (LHsExpr GhcPs, [String])) 73substitute (Subst bind) = transformBracketOld exp . transformBi pat . transformBi typ 74 where 75 exp :: LHsExpr GhcPs -> Maybe (LHsExpr GhcPs) 76 -- Variables. 77 exp (L _ (HsVar _ x)) = lookup (rdrNameStr x) bind 78 -- Operator applications. 79 exp (L loc (OpApp _ lhs (L _ (HsVar _ x)) rhs)) 80 | Just y <- lookup (rdrNameStr x) bind = Just (L loc (OpApp noExtField lhs y rhs)) 81 -- Left sections. 82 exp (L loc (SectionL _ exp (L _ (HsVar _ x)))) 83 | Just y <- lookup (rdrNameStr x) bind = Just (L loc (SectionL noExtField exp y)) 84 -- Right sections. 85 exp (L loc (SectionR _ (L _ (HsVar _ x)) exp)) 86 | Just y <- lookup (rdrNameStr x) bind = Just (L loc (SectionR noExtField y exp)) 87 exp _ = Nothing 88 89 pat :: LPat GhcPs -> LPat GhcPs 90 -- Pattern variables. 91 pat (L _ (VarPat _ x)) 92 | Just y@(L _ HsVar{}) <- lookup (rdrNameStr x) bind = strToPat $ varToStr y 93 pat x = x :: LPat GhcPs 94 95 typ :: LHsType GhcPs -> LHsType GhcPs 96 -- Type variables. 97 typ (L _ (HsTyVar _ _ x)) 98 | Just (L _ (HsAppType _ _ (HsWC _ y))) <- lookup (rdrNameStr x) bind = y 99 typ x = x :: LHsType GhcPs 100 101 102--------------------------------------------------------------------- 103-- UNIFICATION 104 105type NameMatch = Located RdrName -> Located RdrName -> Bool 106 107-- | Unification, obeys the property that if @unify a b = s@, then 108-- @substitute s a = b@. 109unify' :: Data a => NameMatch -> Bool -> a -> a -> Maybe (Subst (LHsExpr GhcPs)) 110unify' nm root x y 111 | Just (x, y) <- cast (x, y) = unifyExp' nm root x y 112 | Just (x, y) <- cast (x, y) = unifyPat' nm x y 113 | Just (x, y) <- cast (x, y) = unifyType' nm x y 114 | Just (x, y) <- cast (x, y) = if (x :: FastString) == y then Just mempty else Nothing 115 | Just (x :: SrcSpan) <- cast x = Just mempty 116 | otherwise = unifyDef' nm x y 117 118unifyDef' :: Data a => NameMatch -> a -> a -> Maybe (Subst (LHsExpr GhcPs)) 119unifyDef' nm x y = fmap mconcat . sequence =<< gzip (unify' nm False) x y 120 121unifyComposed' :: NameMatch 122 -> LHsExpr GhcPs 123 -> LHsExpr GhcPs -> LHsExpr GhcPs -> LHsExpr GhcPs 124 -> Maybe (Subst (LHsExpr GhcPs), Maybe (LHsExpr GhcPs)) 125unifyComposed' nm x1 y11 dot y12 = 126 ((, Just y11) <$> unifyExp' nm False x1 y12) 127 <|> case y12 of 128 (L _ (OpApp _ y121 dot' y122)) | isDot dot' -> 129 unifyComposed' nm x1 (noLoc (OpApp noExtField y11 dot y121)) dot' y122 130 _ -> Nothing 131 132-- unifyExp handles the cases where both x and y are HsApp, or y is OpApp. Otherwise, 133-- delegate to unifyExp'. These are the cases where we potentially need to call 134-- unifyComposed' to handle left composition. 135-- 136-- y is allowed to partially match x (the lhs of the hint), if y is a function application where 137-- the function is a composition of functions. In this case the second component of the result is 138-- the unmatched part of y, which will be attached to the rhs of the hint after substitution. 139-- 140-- Example: 141-- x = head (drop n x) 142-- y = foo . bar . baz . head $ drop 2 xs 143-- result = (Subst [(n, 2), (x, xs)], Just (foo . bar . baz)) 144unifyExp :: NameMatch -> Bool -> LHsExpr GhcPs -> LHsExpr GhcPs -> Maybe (Subst (LHsExpr GhcPs), Maybe (LHsExpr GhcPs)) 145-- Match wildcard operators. 146unifyExp nm root (L _ (OpApp _ lhs1 (L _ (HsVar _ (rdrNameStr -> v))) rhs1)) 147 (L _ (OpApp _ lhs2 (L _ (HsVar _ (rdrNameStr -> op2))) rhs2)) 148 | isUnifyVar v = 149 (, Nothing) . (Subst [(v, strToVar op2)] <>) <$> 150 liftA2 (<>) (unifyExp' nm False lhs1 lhs2) (unifyExp' nm False rhs1 rhs2) 151 152-- Options: match directly, and expand through '.' 153unifyExp nm root x@(L _ (HsApp _ x1 x2)) (L _ (HsApp _ y1 y2)) = 154 ((, Nothing) <$> liftA2 (<>) (unifyExp' nm False x1 y1) (unifyExp' nm False x2 y2)) <|> unifyComposed 155 where 156 -- Unify a function application where the function is a composition of functions. 157 unifyComposed 158 | (L _ (OpApp _ y11 dot y12)) <- fromParen y1, isDot dot = 159 if not root then 160 -- Attempt #1: rewrite '(fun1 . fun2) arg' as 'fun1 (fun2 arg)', and unify it with 'x'. 161 -- The guard ensures that you don't get duplicate matches because the matching engine 162 -- auto-generates hints in dot-form. 163 (, Nothing) <$> unifyExp' nm root x (noLoc (HsApp noExtField y11 (noLoc (HsApp noExtField y12 y2)))) 164 else do 165 -- Attempt #2: rewrite '(fun1 . fun2 ... funn) arg' as 'fun1 $ (fun2 ... funn) arg', 166 -- 'fun1 . fun2 $ (fun3 ... funn) arg', 'fun1 . fun2 . fun3 $ (fun4 ... funn) arg', 167 -- and so on, unify the rhs of '$' with 'x', and store the lhs of '$' into 'extra'. 168 -- You can only add to extra if you are at the root (otherwise 'extra' has nowhere to go). 169 rhs <- unifyExp' nm False x2 y2 170 (lhs, extra) <- unifyComposed' nm x1 y11 dot y12 171 pure (lhs <> rhs, extra) 172 | otherwise = Nothing 173 174-- Options: match directly, then expand through '$', then desugar infix. 175unifyExp nm root x (L _ (OpApp _ lhs2 op2@(L _ (HsVar _ op2')) rhs2)) 176 | (L _ (OpApp _ lhs1 op1@(L _ (HsVar _ op1')) rhs1)) <- x = 177 guard (nm op1' op2') >> (, Nothing) <$> liftA2 (<>) (unifyExp' nm False lhs1 lhs2) (unifyExp' nm False rhs1 rhs2) 178 | isDol op2 = unifyExp nm root x $ noLoc (HsApp noExtField lhs2 rhs2) 179 | isAmp op2 = unifyExp nm root x $ noLoc (HsApp noExtField rhs2 lhs2) 180 | otherwise = unifyExp nm root x $ noLoc (HsApp noExtField (noLoc (HsApp noExtField op2 (addPar lhs2))) (addPar rhs2)) 181 where 182 -- add parens around when desugaring the expression, if necessary 183 addPar :: LHsExpr GhcPs -> LHsExpr GhcPs 184 addPar x = if isAtom x then x else addParen x 185 186unifyExp nm root x y = (, Nothing) <$> unifyExp' nm root x y 187 188isAmp :: LHsExpr GhcPs -> Bool 189isAmp (L _ (HsVar _ x)) = rdrNameStr x == "&" 190isAmp _ = False 191 192-- | If we "throw away" the extra than we have no where to put it, and the substitution is wrong 193noExtra :: Maybe (Subst (LHsExpr GhcPs), Maybe (LHsExpr GhcPs)) -> Maybe (Subst (LHsExpr GhcPs)) 194noExtra (Just (x, Nothing)) = Just x 195noExtra _ = Nothing 196 197-- App/InfixApp are analysed specially for performance reasons. If 198-- 'root = True', this is the outside of the expr. Do not expand out a 199-- dot at the root, since otherwise you get two matches because of 200-- 'readRule' (Bug #570). 201unifyExp' :: NameMatch -> Bool -> LHsExpr GhcPs -> LHsExpr GhcPs -> Maybe (Subst (LHsExpr GhcPs)) 202-- Don't subsitute for type apps, since no one writes rules imagining 203-- they exist. 204unifyExp' nm root (L _ (HsVar _ (rdrNameStr -> v))) y | isUnifyVar v, not $ isTypeApp y = Just $ Subst [(v, y)] 205unifyExp' nm root (L _ (HsVar _ x)) (L _ (HsVar _ y)) | nm x y = Just mempty 206 207-- Brackets are not added when expanding '$' in user code, so tolerate 208-- them in the match even if they aren't in the user code. 209-- Also, allow the user to put in more brackets than they strictly need (e.g. with infix). 210unifyExp' nm root x y | not root, isJust x2 || isJust y2 = unifyExp' nm root (fromMaybe x x2) (fromMaybe y y2) 211 where 212 -- Make sure we deal with the weird brackets that can't be removed around sections 213 x2 = remParen x 214 y2 = remParen y 215 216unifyExp' nm root x@(L _ (OpApp _ lhs1 (L _ (HsVar _ (rdrNameStr -> v))) rhs1)) 217 y@(L _ (OpApp _ lhs2 (L _ (HsVar _ op2)) rhs2)) = 218 noExtra $ unifyExp nm root x y 219unifyExp' nm root (L _ (SectionL _ exp1 (L _ (HsVar _ (rdrNameStr -> v))))) 220 (L _ (SectionL _ exp2 (L _ (HsVar _ (rdrNameStr -> op2))))) 221 | isUnifyVar v = (Subst [(v, strToVar op2)] <>) <$> unifyExp' nm False exp1 exp2 222unifyExp' nm root (L _ (SectionR _ (L _ (HsVar _ (rdrNameStr -> v))) exp1)) 223 (L _ (SectionR _ (L _ (HsVar _ (rdrNameStr -> op2))) exp2)) 224 | isUnifyVar v = (Subst [(v, strToVar op2)] <>) <$> unifyExp' nm False exp1 exp2 225 226unifyExp' nm root x@(L _ (HsApp _ x1 x2)) y@(L _ (HsApp _ y1 y2)) = 227 noExtra $ unifyExp nm root x y 228 229unifyExp' nm root x y@(L _ (OpApp _ lhs2 op2@(L _ (HsVar _ op2')) rhs2)) = 230 noExtra $ unifyExp nm root x y 231 232unifyExp' nm root (L _ (HsBracket _ (VarBr _ b0 (occNameStr -> v1)))) 233 (L _ (HsBracket _ (VarBr _ b1 (occNameStr -> v2)))) 234 | b0 == b1 && isUnifyVar v1 = Just (Subst [(v1, strToVar v2)]) 235 236unifyExp' nm root x y | isOther x, isOther y = unifyDef' nm x y 237 where 238 -- Types that are not already handled in unify. 239 {-# INLINE isOther #-} 240 isOther :: LHsExpr GhcPs -> Bool 241 isOther (L _ HsVar{}) = False 242 isOther (L _ HsApp{}) = False 243 isOther (L _ OpApp{}) = False 244 isOther _ = True 245 246unifyExp' _ _ _ _ = Nothing 247 248 249unifyPat' :: NameMatch -> LPat GhcPs -> LPat GhcPs -> Maybe (Subst (LHsExpr GhcPs)) 250unifyPat' nm (L _ (VarPat _ x)) (L _ (VarPat _ y)) = 251 Just $ Subst [(rdrNameStr x, strToVar(rdrNameStr y))] 252unifyPat' nm (L _ (VarPat _ x)) (L _ (WildPat _)) = 253 let s = rdrNameStr x in Just $ Subst [(s, strToVar("_" ++ s))] 254unifyPat' nm (L _ (ConPat _ x _)) (L _ (ConPat _ y _)) | rdrNameStr x /= rdrNameStr y = 255 Nothing 256unifyPat' nm x y = 257 unifyDef' nm x y 258 259unifyType' :: NameMatch -> LHsType GhcPs -> LHsType GhcPs -> Maybe (Subst (LHsExpr GhcPs)) 260unifyType' nm (L loc (HsTyVar _ _ x)) y = 261 let wc = HsWC noExtField y :: LHsWcType (NoGhcTc GhcPs) 262 unused = strToVar "__unused__" :: LHsExpr GhcPs 263 appType = L loc (HsAppType noExtField unused wc) :: LHsExpr GhcPs 264 in Just $ Subst [(rdrNameStr x, appType)] 265unifyType' nm x y = unifyDef' nm x y 266