@@ -12,8 +12,9 @@ where
1212
1313import Data.Function (on)
1414import qualified Data.List as L
15+import Data.Map.Strict (Map)
16+import qualified Data.Map.Strict as M
1517import Data.Maybe (fromMaybe, mapMaybe)
16−import Data.Ord (Down (Down), comparing)
1718import GHC
1819import OccName (mkVarOcc)
1920import Ormolu.Utils (unSrcSpan)
@@ -56,7 +57,7 @@ reassociateOpTree getOpName opTree =
5657reassociateOpTreeWith ::
5758 forall ty op.
5859 -- | Fixity map for operators
59− [(RdrName, Fixity)] ->
60+ Map RdrName Fixity ->
6061 -- | How to get the name of an operator
6162 (op -> Maybe RdrName) ->
6263 -- | Original 'OpTree'
@@ -68,12 +69,12 @@ reassociateOpTreeWith fixityMap getOpName = go
6869 fixityOf :: op -> Fixity
6970 fixityOf op = fromMaybe defaultFixity $ do
7071 opName <- getOpName op
71− lookup opName fixityMap
72+ M.lookup opName fixityMap
7273 -- Here, left branch is already associated and the root alongside with
7374 -- the right branch is right-associated. This function picks up one item
7475 -- from the right and inserts it correctly to the left.
7576 --
76− -- Also, we are using the 'compareFixity' function which returns if the
77+ -- Also, we are using the 'compareFixity' function which tells if the
7778 -- expression should associate to right.
7879 go :: OpTree ty op -> OpTree ty op
7980 -- base cases
@@ -93,6 +94,16 @@ reassociateOpTreeWith fixityMap getOpName = go
9394 then go $ OpBranch (OpBranch l op (go $ OpBranch r op' l')) op'' r'
9495 else go $ OpBranch (OpBranch (OpBranch l op r) op' l') op'' r'
9596
97+-- | A score assigned to an operator.
98+data Score
99+ = -- | The operator was placed at the beginning of a line
100+ AtBeginning Int
101+ | -- | The operator was placed at the end of a line
102+ AtEnd
103+ | -- | The operator was placed in between arguments on a single line
104+ InBetween
105+ deriving (Eq, Ord)
106+
96107-- | Build a map of inferred 'Fixity's from an 'OpTree'.
97108buildFixityMap ::
98109 forall ty op.
@@ -101,23 +112,26 @@ buildFixityMap ::
101112 -- | Operator tree
102113 OpTree (Located ty) (Located op) ->
103114 -- | Fixity map
104− [(RdrName, Fixity)]
115+ Map RdrName Fixity
105116buildFixityMap getOpName opTree =
106− concatMap (\(i, ns) -> map (\(n, _) -> (n, fixity i InfixL)) ns)
107− . zip [0 ..]
108− . L.groupBy (doubleWithinEps 0.00001 `on` snd)
109− . (overrides ++)
110− . modeScores
117+ addOverrides
118+ . M.fromList
119+ . concatMap (\(i, ns) -> map (\(n, _) -> (n, fixity i InfixL)) ns)
120+ . zip [1 ..]
121+ . L.groupBy ((==) `on` snd)
122+ . selectScores
111123 $ score opTree
112124 where
113− -- Add a special case for ($), since it is pretty unlikely for someone
114− -- to override it.
115− overrides :: [(RdrName, Double)]
116− overrides =
117− [ (mkRdrUnqual $ mkVarOcc "$", -1)
118− ]
119− -- Assign scores to operators based on their location in the source.
120− score :: OpTree (Located ty) (Located op) -> [(RdrName, Double)]
125+ addOverrides :: Map RdrName Fixity -> Map RdrName Fixity
126+ addOverrides m =
127+ let mk k v = (mkRdrUnqual (mkVarOcc k), fixity v InfixL)
128+ in M.fromList
129+ [ mk "$" 0,
130+ mk "." 9
131+ ]
132+ `M.union` m
133+ fixity = Fixity NoSourceText
134+ score :: OpTree (Located ty) (Located op) -> [(RdrName, Score)]
121135 score (OpNode _) = []
122136 score (OpBranch l o r) = fromMaybe (score r) $ do
123137 -- If we fail to get any of these, 'defaultFixity' will be used by
@@ -129,46 +143,25 @@ buildFixityMap getOpName opTree =
129143 oc <- srcSpanStartCol <$> unSrcSpan (getLoc o) -- operator column
130144 opName <- getOpName (unLoc o)
131145 let s
132− | le < ob =
133− -- if the operator is in the beginning of a line, assign
134− -- a score relative to its column within range [0, 1).
135− fromIntegral oc / fromIntegral (maxCol + 1)
136− | oe < rb =
137− -- if the operator is in the end of the line, assign the
138− -- score 1.
139− 1
140− | otherwise =
141− 2 -- otherwise, assign a high score.
146+ | le < ob = AtBeginning oc
147+ | oe < rb = AtEnd
148+ | otherwise = InBetween
142149 return $ (opName, s) : score r
143− -- Pick the most common score per 'RdrName'.
144− modeScores :: [(RdrName, Double)] -> [(RdrName, Double)]
145− modeScores =
150+ selectScores :: [(RdrName, Score)] -> [(RdrName, Score)]
151+ selectScores =
146152 L.sortOn snd
147153 . mapMaybe
148154 ( \case
149155 [] -> Nothing
150− xs@((n, _) : _) -> Just (n, mode $ map snd xs)
156+ xs@((n, _) : _) -> Just (n, selectScore $ map snd xs)
151157 )
152158 . L.groupBy ((==) `on` fst)
153159 . L.sort
154− -- Return the most common number, leaning to the smaller
155− -- one in case of a tie.
156− mode :: [Double] -> Double
157− mode =
158− head
159− . L.minimumBy (comparing (Down . length))
160− . L.groupBy (doubleWithinEps 0.0001)
161− . L.sort
162− -- The start column of the rightmost operator.
163− maxCol = go opTree
164− where
165− go (OpNode (L _ _)) = 0
166− go (OpBranch l (L o _) r) =
167− maximum
168− [ go l,
169− maybe 0 srcSpanStartCol (unSrcSpan o),
170− go r
171− ]
160+ selectScore :: [Score] -> Score
161+ selectScore xs =
162+ case filter (/= InBetween) xs of
163+ [] -> InBetween
164+ xs' -> maximum xs'
172165
173166----------------------------------------------------------------------------
174167-- Helpers
@@ -182,9 +175,3 @@ normalizeOpTree (OpBranch (OpNode l) lop r) =
182175 OpBranch (OpNode l) lop (normalizeOpTree r)
183176normalizeOpTree (OpBranch (OpBranch l' lop' r') lop r) =
184177 normalizeOpTree (OpBranch l' lop' (OpBranch r' lop r))
185−
186−fixity :: Int -> FixityDirection -> Fixity
187−fixity = Fixity NoSourceText
188−
189−doubleWithinEps :: Double -> Double -> Double -> Bool
190−doubleWithinEps eps a b = abs (a - b) < eps