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