Skip to content

Commit 88b8e82

Browse files
zliu41zliu41
authored
Add brackets if needed for hint "Avoid lambda" (#1634)
1 parent 4620d86 commit 88b8e82

File tree

2 files changed

+109
-93
lines changed

2 files changed

+109
-93
lines changed

‎src/GHC/Util/HsExpr.hs

+104-92
Original file line numberDiff line numberDiff line change
@@ -141,7 +141,7 @@ niceDotApp a b = dotApp a b
141141

142142
-- Generate a lambda expression but prettier if possible.
143143
niceLambda :: [String] -> LHsExpr GhcPs -> LHsExpr GhcPs
144-
niceLambda ss e = fst (niceLambdaR ss e)-- We don't support refactorings yet.
144+
niceLambda ss e = fst (niceLambdaR Nothing ss e)-- We don't support refactorings yet.
145145

146146
allowRightSection :: String -> Bool
147147
allowRightSection x = x `notElem` ["-","#"]
@@ -150,99 +150,111 @@ allowLeftSection x = x /= "#"
150150

151151
-- Implementation. Try to produce special forms (e.g. sections,
152152
-- compositions) where we can.
153-
niceLambdaR :: [String]
154-
-> LHsExpr GhcPs
153+
niceLambdaR :: Maybe (LHsExpr GhcPs) -- parent expression
154+
-> [String]
155+
-> LHsExpr GhcPs -- the expression being processed
155156
-> (LHsExpr GhcPs, R.SrcSpan -> [Refactoring R.SrcSpan])
156-
-- Rewrite @\ -> e@ as @e@
157-
-- These are encountered as recursive calls.
158-
niceLambdaR xs (SimpleLambda [] x) = niceLambdaR xs x
159-
160-
-- Rewrite @\xs -> (e)@ as @\xs -> e@.
161-
niceLambdaR xs (L _ (HsPar _ x)) = niceLambdaR xs x
162-
163-
-- @\vs v -> ($) e v@ ==> @\vs -> e@
164-
-- @\vs v -> e $ v@ ==> @\vs -> e@
165-
niceLambdaR (unsnoc -> Just (vs, v)) (view -> App2 f e (view -> Var_ v'))
166-
| isDol f
167-
, v == v'
168-
, vars e `disjoint` [v]
169-
= niceLambdaR vs e
170-
171-
-- @\v -> thing + v@ ==> @\v -> (thing +)@ (heuristic: @v@ must be a single
172-
-- lexeme, or it all gets too complex)
173-
niceLambdaR [v] (L _ (OpApp _ e f (view -> Var_ v')))
174-
| isLexeme e
175-
, v == v'
176-
, vars e `disjoint` [v]
177-
, L _ (HsVar _ (L _ fname)) <- f
178-
, isSymOcc $ rdrNameOcc fname
179-
= let res = nlHsPar $ noLocA $ SectionL noExtField e f
180-
in (res, \s -> [Replace Expr s [] (unsafePrettyPrint res)])
181-
182-
-- @\vs v -> f x v@ ==> @\vs -> f x@
183-
niceLambdaR (unsnoc -> Just (vs, v)) (L _ (HsApp _ f (view -> Var_ v')))
184-
| v == v'
185-
, vars f `disjoint` [v]
186-
= niceLambdaR vs f
187-
188-
-- @\vs v -> (v `f`)@ ==> @\vs -> f@
189-
niceLambdaR (unsnoc -> Just (vs, v)) (L _ (SectionL _ (view -> Var_ v') f))
190-
| v == v' = niceLambdaR vs f
191-
192-
-- Strip one variable pattern from the end of a lambdas match, and place it in our list of factoring variables.
193-
niceLambdaR xs (SimpleLambda ((view -> PVar_ v):vs) x)
194-
| v `notElem` xs = niceLambdaR (xs++[v]) $ lambda vs x
195-
196-
-- Rewrite @\x -> x + a@ as @(+ a)@ (heuristic: @a@ must be a single
197-
-- lexeme, or it all gets too complex).
198-
niceLambdaR [x] (view -> App2 op@(L _ (HsVar _ (L _ tag))) l r)
199-
| isLexeme r, view l == Var_ x, x `notElem` vars r, allowRightSection (occNameStr tag) =
200-
let e = rebracket1 $ addParen (noLocA $ SectionR noExtField op r)
201-
in (e, \s -> [Replace Expr s [] (unsafePrettyPrint e)])
202-
-- Rewrite (1) @\x -> f (b x)@ as @f . b@, (2) @\x -> f $ b x@ as @f . b@.
203-
niceLambdaR [x] y
204-
| Just (z, subts) <- factor y, x `notElem` vars z = (z, \s -> [mkRefact subts s])
157+
niceLambdaR parent = go
205158
where
206-
-- Factor the expression with respect to x.
207-
factor :: LHsExpr GhcPs -> Maybe (LHsExpr GhcPs, [LHsExpr GhcPs])
208-
factor (L _ (HsApp _ ini lst)) | view lst == Var_ x = Just (ini, [ini])
209-
factor (L _ (HsApp _ ini lst)) | Just (z, ss) <- factor lst
210-
= let r = niceDotApp ini z
211-
in if astEq r z then Just (r, ss) else Just (r, ini : ss)
212-
factor (L _ (OpApp _ y op (factor -> Just (z, ss))))| isDol op
213-
= let r = niceDotApp y z
214-
in if astEq r z then Just (r, ss) else Just (r, y : ss)
215-
factor (L _ (HsPar _ y@(L _ HsApp{}))) = factor y
216-
factor _ = Nothing
217-
mkRefact :: [LHsExpr GhcPs] -> R.SrcSpan -> Refactoring R.SrcSpan
218-
mkRefact subts s =
219-
let tempSubts = zipWith (\a b -> (a, toSSA b)) substVars subts
220-
template = dotApps (map (strToVar . fst) tempSubts)
221-
in Replace Expr s tempSubts (unsafePrettyPrint template)
222-
-- Rewrite @\x y -> x + y@ as @(+)@.
223-
niceLambdaR [x,y] (L _ (OpApp _ (view -> Var_ x1) op@(L _ HsVar {}) (view -> Var_ y1)))
224-
| x == x1, y == y1, vars op `disjoint` [x, y] = (op, \s -> [Replace Expr s [] (unsafePrettyPrint op)])
225-
-- Rewrite @\x y -> f y x@ as @flip f@.
226-
niceLambdaR [x, y] (view -> App2 op (view -> Var_ y1) (view -> Var_ x1))
227-
| x == x1, y == y1, vars op `disjoint` [x, y] =
228-
( gen op
229-
, \s -> [Replace Expr s [("x", toSSA op)] (unsafePrettyPrint $ gen (strToVar "x"))]
230-
)
231-
where
232-
gen :: LHsExpr GhcPs -> LHsExpr GhcPs
233-
gen = noLocA . HsApp noExtField (strToVar "flip")
234-
. if isAtom op then id else addParen
235-
236-
-- We're done factoring, but have no variables left, so we shouldn't make a lambda.
237-
-- @\ -> e@ ==> @e@
238-
niceLambdaR [] e = (e, \s -> [Replace Expr s [("a", toSSA e)] "a"])
239-
-- Base case. Just a good old fashioned lambda.
240-
niceLambdaR ss e =
241-
let grhs = noLocA $ GRHS noAnn [] e :: LGRHS GhcPs (LHsExpr GhcPs)
242-
grhss = GRHSs {grhssExt = emptyComments, grhssGRHSs=[grhs], grhssLocalBinds=EmptyLocalBinds noExtField}
243-
match = noLocA $ Match {m_ext=noExtField, m_ctxt=LamAlt LamSingle, m_pats=noLocA $ map strToPat ss, m_grhss=grhss} :: LMatch GhcPs (LHsExpr GhcPs)
244-
matchGroup = MG {mg_ext=Generated OtherExpansion SkipPmc, mg_alts=noLocA [match]}
245-
in (noLocA $ HsLam noAnn LamSingle matchGroup, const [])
159+
-- Rewrite @\ -> e@ as @e@
160+
-- These are encountered as recursive calls.
161+
go xs (SimpleLambda [] x) = go xs x
162+
163+
-- Rewrite @\xs -> (e)@ as @\xs -> e@.
164+
go xs (L _ (HsPar _ x)) = go xs x
165+
166+
-- @\vs v -> ($) e v@ ==> @\vs -> e@
167+
-- @\vs v -> e $ v@ ==> @\vs -> e@
168+
go (unsnoc -> Just (vs, v)) (view -> App2 f e (view -> Var_ v'))
169+
| isDol f
170+
, v == v'
171+
, vars e `disjoint` [v]
172+
= go vs e
173+
174+
-- @\v -> thing + v@ ==> @\v -> (thing +)@ (heuristic: @v@ must be a single
175+
-- lexeme, or it all gets too complex)
176+
go [v] (L _ (OpApp _ e f (view -> Var_ v')))
177+
| isLexeme e
178+
, v == v'
179+
, vars e `disjoint` [v]
180+
, L _ (HsVar _ (L _ fname)) <- f
181+
, isSymOcc $ rdrNameOcc fname
182+
= let res = nlHsPar $ noLocA $ SectionL noExtField e f
183+
in (res, \s -> [Replace Expr s [] (unsafePrettyPrint res)])
184+
185+
-- @\vs v -> f x v@ ==> @\vs -> f x@
186+
go (unsnoc -> Just (vs, v)) (L _ (HsApp _ f (view -> Var_ v')))
187+
| v == v'
188+
, vars f `disjoint` [v]
189+
= go vs f
190+
191+
-- @\vs v -> (v `f`)@ ==> @\vs -> f@
192+
go (unsnoc -> Just (vs, v)) (L _ (SectionL _ (view -> Var_ v') f))
193+
| v == v' = go vs f
194+
195+
-- Strip one variable pattern from the end of a lambdas match, and place it in our list of factoring variables.
196+
go xs (SimpleLambda ((view -> PVar_ v):vs) x)
197+
| v `notElem` xs = go (xs++[v]) $ lambda vs x
198+
199+
-- Rewrite @\x -> x + a@ as @(+ a)@ (heuristic: @a@ must be a single
200+
-- lexeme, or it all gets too complex).
201+
go [x] (view -> App2 op@(L _ (HsVar _ (L _ tag))) l r)
202+
| isLexeme r, view l == Var_ x, x `notElem` vars r, allowRightSection (occNameStr tag) =
203+
let e = rebracket1 $ addParen (noLocA $ SectionR noExtField op r)
204+
in (e, \s -> [Replace Expr s [] (unsafePrettyPrint e)])
205+
-- Rewrite (1) @\x -> f (b x)@ as @f . b@, (2) @\x -> f $ b x@ as @f . b@.
206+
go [x] y
207+
| Just (z, subts) <- factor y, x `notElem` vars z = (z, \s -> [mkRefact subts s])
208+
where
209+
-- Factor the expression with respect to x.
210+
factor :: LHsExpr GhcPs -> Maybe (LHsExpr GhcPs, [LHsExpr GhcPs])
211+
factor (L _ (HsApp _ ini lst)) | view lst == Var_ x = Just (ini, [ini])
212+
factor (L _ (HsApp _ ini lst)) | Just (z, ss) <- factor lst
213+
= let r = niceDotApp ini z
214+
in if astEq r z then Just (r, ss) else Just (r, ini : ss)
215+
factor (L _ (OpApp _ y op (factor -> Just (z, ss))))| isDol op
216+
= let r = niceDotApp y z
217+
in if astEq r z then Just (r, ss) else Just (r, y : ss)
218+
factor (L _ (HsPar _ y@(L _ HsApp{}))) = factor y
219+
factor _ = Nothing
220+
mkRefact :: [LHsExpr GhcPs] -> R.SrcSpan -> Refactoring R.SrcSpan
221+
mkRefact subts s =
222+
let tempSubts = zipWith (\a b -> (a, toSSA b)) substVars subts
223+
template = dotApps (map (strToVar . fst) tempSubts)
224+
in Replace Expr s tempSubts (unsafePrettyPrint template)
225+
-- Rewrite @\x y -> x + y@ as @(+)@.
226+
go [x,y] (L _ (OpApp _ (view -> Var_ x1) op@(L _ HsVar {}) (view -> Var_ y1)))
227+
| x == x1, y == y1, vars op `disjoint` [x, y] = (op, \s -> [Replace Expr s [] (unsafePrettyPrint op)])
228+
-- Rewrite @\x y -> f y x@ as @flip f@.
229+
go [x, y] (view -> App2 op (view -> Var_ y1) (view -> Var_ x1))
230+
| x == x1, y == y1, vars op `disjoint` [x, y] =
231+
( gen op
232+
, \s -> [Replace Expr s [("x", toSSA op)] (unsafePrettyPrint $ gen (strToVar "x"))]
233+
)
234+
where
235+
gen :: LHsExpr GhcPs -> LHsExpr GhcPs
236+
gen = noLocA . HsApp noExtField (strToVar "flip")
237+
. if isAtom op then id else addParen
238+
239+
-- We're done factoring, but have no variables left, so we shouldn't make a lambda.
240+
-- @\ -> e@ ==> @e@
241+
go [] e =
242+
let -- Add brackets if needed, primarily for handling BlockArguments.
243+
-- e.g., parent = `f \x -> g 3 x`; e = `g 3`.
244+
-- Brackets should be placed around `e` to produce `f (g 3)` instead of `f g 3`.
245+
addBrackets = case parent of
246+
Just p -> isApp p && not (isVar e)
247+
Nothing -> False
248+
e' = if addBrackets then mkHsPar e else e
249+
tpl = if addBrackets then "(a)" else "a"
250+
in (e', \s -> [Replace Expr s [("a", toSSA e)] tpl])
251+
-- Base case. Just a good old fashioned lambda.
252+
go ss e =
253+
let grhs = noLocA $ GRHS noAnn [] e :: LGRHS GhcPs (LHsExpr GhcPs)
254+
grhss = GRHSs {grhssExt = emptyComments, grhssGRHSs=[grhs], grhssLocalBinds=EmptyLocalBinds noExtField}
255+
match = noLocA $ Match {m_ext=noExtField, m_ctxt=LamAlt LamSingle, m_pats=noLocA $ map strToPat ss, m_grhss=grhss} :: LMatch GhcPs (LHsExpr GhcPs)
256+
matchGroup = MG {mg_ext=Generated OtherExpansion SkipPmc, mg_alts=noLocA [match]}
257+
in (noLocA $ HsLam noAnn LamSingle matchGroup, const [])
246258

247259

248260
-- 'case' and 'if' expressions have branches, nothing else does (this

‎src/Hint/Lambda.hs

+5-1
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,10 @@ f = foo (\y -> g x . h $ y) -- g x . h
3939
f = foo (\y -> g x . h $ y) -- @Message Avoid lambda
4040
f = foo ((*) x) -- (x *)
4141
f = (*) x
42+
f = g \x -> h 3 x -- (h 3)
43+
f = g (\x -> h 3 x) -- h 3
44+
f = g \x -> (`h` 3) x -- (`h` 3)
45+
f = g \x -> h x -- h
4246
f = foo (flip op x) -- (`op` x)
4347
f = foo (flip op x) -- @Message Use section
4448
f = foo (flip x y) -- (`x` y)
@@ -217,7 +221,7 @@ lambdaExp _ o@(L _ (HsPar _ (view -> App2 (view -> Var_ "flip") origf@(view -> R
217221

218222
lambdaExp p o@(L _ (HsLam _ LamSingle _))
219223
| not $ any isOpApp p
220-
, (res, refact) <- niceLambdaR [] o
224+
, (res, refact) <- niceLambdaR p [] o
221225
, not $ isLambda res
222226
, not $ any isQuasiQuoteExpr $ universe res
223227
, not $ "runST" `Set.member` Set.map occNameString (freeVars o)

0 commit comments

Comments
 (0)