Handle operator chains better

I implemented a custom logic where we assign a score to every occurance of
an operator based on their location, and the average of that score determine
the fixity of the operator.

As you can imagine, the solution is a bit brittle; and it is easy to mislead
it if you knowingly craft an input, but it gave acceptable results for every
code snippet I found online. And since it returns the same AST no matter how
we infer the fixities, it is not the end of the world if we infer something
incorrectly.

The code is not really optimised, and I think it has quadratic time
complexity. Notably, we use opTreeLoc function quite often and it traverses
the whole tree every time. Memoizing that on the OpBranch constructor would
make formatting files with reeeally long operator chains a lot faster. We
can do this once we decide to optimize for speed.

data/examples/declaration/type-synonyms/multi-line-out.hs

@@ -14,8 +14,5 @@ type Foo
 
 type API
   = "route1" :> ApiRoute1
-      :<|> "route2"
-      :> ApiRoute2 -- comment here
-      :<|> OmitDocs
-      :> "i"
-      :> ASomething API
+      :<|> "route2" :> ApiRoute2 -- comment here
+      :<|> OmitDocs :> "i" :> ASomething API

data/examples/declaration/value/function/infix/lenses-out.hs

@@ -1,11 +1,20 @@
 lenses =
   Just $ M.fromList
-    $ "type"
-    .= ("user.connection" :: Text)
-    # "connection"
-    .= uc
-    # "user"
-    .= case name of
-      Just n -> Just $ object ["name" .= n]
-      Nothing -> Nothing
-    # []
+    $ "type" .= ("user.connection" :: Text)
+      # "connection" .= uc
+      # "user" .= case name of
+        Just n -> Just $ object ["name" .= n]
+        Nothing -> Nothing
+      # []
+
+foo =
+  a
+    & b .~ 2
+    & c .~ 3
+
+wreq =
+  let opts =
+        defaults & auth ?~ awsAuth AWSv4 "key" "secret"
+          & header "Accept" .~ ["application/json"]
+          & header "Runscope-Bucket-Auth" .~ ["1example-1111-4yyyy-zzzz-xxxxxxxx"]
+   in getWith opts

data/examples/declaration/value/function/infix/lenses.hs

@@ -5,3 +5,12 @@ lenses = Just $ M.fromList
       Just  n -> Just $ object ["name" .= n]
       Nothing -> Nothing
   # []
+
+foo = a
+  & b .~ 2 & c .~ 3
+
+wreq =
+  let opts = defaults & auth ?~ awsAuth AWSv4 "key" "secret"
+                          & header "Accept" .~ ["application/json"]
+                          & header "Runscope-Bucket-Auth" .~ ["1example-1111-4yyyy-zzzz-xxxxxxxx"]
+  in  getWith opts

data/examples/declaration/value/function/list-comprehensions-out.hs

@@ -10,8 +10,7 @@ barbaz x y z w =
       any even [a, b],
       c <-
         z
-          * z
-          ^ 2, -- Bar baz
+          * z ^ 2, -- Bar baz
       d <-
         w
           + w, -- Baz bar

default.nix

@@ -53,6 +53,7 @@ in {
       "servant"
       "servant-server"
       "tensorflow"
+      "text_1_2_4_0"
       "tls"
       "yesod-core"
 

ormolu.cabal

@@ -106,9 +106,10 @@ library
                     , Ormolu.Printer.Meat.Declaration.Value
                     , Ormolu.Printer.Meat.Declaration.Warning
                     , Ormolu.Printer.Meat.ImportExport
-                    , Ormolu.Printer.Meat.Pragma
                     , Ormolu.Printer.Meat.Module
+                    , Ormolu.Printer.Meat.Pragma
                     , Ormolu.Printer.Meat.Type
+                    , Ormolu.Printer.Operators
                     , Ormolu.Printer.SpanStream
                     , Ormolu.Utils
   if flag(dev)

src/Ormolu/Printer/Meat/Declaration/Value.hs

@@ -24,6 +24,7 @@ import Ormolu.Printer.Combinators
 import Ormolu.Printer.Meat.Common
 import Ormolu.Printer.Meat.Declaration.Signature
 import Ormolu.Printer.Meat.Type
+import Ormolu.Printer.Operators
 import Ormolu.Utils
 import Outputable (Outputable (..))
 import RdrName (rdrNameOcc)
@@ -507,27 +508,11 @@ p_hsExpr = \case
       txt "@"
       located (hswc_body a) p_hsType
   OpApp NoExt x op y -> do
-    -- NOTE If the beginning of the first argument and the second argument
-    -- are on the same line, and the second argument has a hanging form, use
-    -- hanging placement.
-    let placement =
-          if isOneLineSpan
-               (mkSrcSpan (srcSpanStart (getLoc x)) (srcSpanStart (getLoc y)))
-            then exprPlacement (unLoc y)
-            else Normal
-        opWrapper = case unLoc op of
-          EWildPat NoExt -> backticks
-          _ -> id
-    ub <- vlayout
-      (return useBraces)
-      (return $ case placement of
-         Hanging -> useBraces
-         Normal -> dontUseBraces)
-    ub $ located x p_hsExpr
-    placeHanging placement $ do
-      located op (opWrapper . p_hsExpr)
-      space
-      located y p_hsExpr
+    let opTree = OpBranch (exprOpTree x) op (exprOpTree y)
+        getOpName = \case
+          HsVar NoExt (L _ a) -> Just a
+          _ -> Nothing
+    p_exprOpTree (reassociateOpTree getOpName opTree)
   NegApp NoExt e _ -> do
     txt "-"
     space
@@ -1042,6 +1027,39 @@ withGuards = any (checkOne . unLoc)
     checkOne (GRHS NoExt [] _) = False
     checkOne _ = True
 
+exprOpTree :: LHsExpr GhcPs -> OpTree (LHsExpr GhcPs) (LHsExpr GhcPs)
+exprOpTree (L _ (OpApp NoExt x op y)) = OpBranch (exprOpTree x) op (exprOpTree y)
+exprOpTree n = OpNode n
+
+p_exprOpTree :: OpTree (LHsExpr GhcPs) (LHsExpr GhcPs) -> R ()
+p_exprOpTree (OpNode x) = located x p_hsExpr
+p_exprOpTree (OpBranch x op y) = do
+  -- NOTE If the beginning of the first argument and the second argument
+  -- are on the same line, and the second argument has a hanging form, use
+  -- hanging placement.
+  let placement =
+        if isOneLineSpan
+             (mkSrcSpan (srcSpanStart (opTreeLoc x)) (srcSpanStart (opTreeLoc y)))
+          then case y of
+                 OpNode (L _ n) -> exprPlacement n
+                 _ -> Normal
+          else Normal
+      opWrapper = case unLoc op of
+        EWildPat NoExt -> backticks
+        _ -> id
+  ub <- vlayout
+    (return useBraces)
+    (return $ case placement of
+       Hanging -> useBraces
+       Normal -> dontUseBraces)
+  switchLayout [opTreeLoc x] $
+    ub $ p_exprOpTree x
+  placeHanging placement $ do
+    located op (opWrapper . p_hsExpr)
+    space
+    switchLayout [opTreeLoc y] $
+      p_exprOpTree y
+
 -- | Get annotations for the enclosing element.
 
 getEnclosingAnns :: R [AnnKeywordId]

src/Ormolu/Printer/Meat/Type.hs

@@ -16,8 +16,8 @@ where
 import GHC
 import Ormolu.Printer.Combinators
 import Ormolu.Printer.Meat.Common
+import Ormolu.Printer.Operators
 import Ormolu.Utils
-import SrcLoc (combineSrcSpans)
 import {-# SOURCE #-} Ormolu.Printer.Meat.Declaration.Value (p_hsSplice)
 
 p_hsType :: HsType GhcPs -> R ()
@@ -72,25 +72,8 @@ p_hsType = \case
     parensHash . sitcc $
       sep (txt "|" >> breakpoint) (sitcc . located' p_hsType) xs
   HsOpTy NoExt x op y -> sitcc $ do
-    -- In the AST, type operators are right-associative instead of left-associative
-    -- like value level operators. This makes similar constructs look inconsistent.
-    -- Here, we shake the AST to convert right-associative tree to a left-associative
-    -- one.
-    case unLoc y of
-      HsOpTy NoExt x' op' y' ->
-        p_hsType $
-          HsOpTy
-            NoExt
-            (L (combineSrcSpans (getLoc x) (getLoc x')) (HsOpTy NoExt x op x'))
-            op'
-            y'
-      _ -> do
-        located x p_hsType
-        breakpoint
-        inci $ do
-          p_rdrName op
-          space
-          located y p_hsType
+    let opTree = OpBranch (tyOpTree x) op (tyOpTree y)
+     in p_tyOpTree (reassociateOpTree Just opTree)
   HsParTy NoExt (L _ t@HsKindSig {}) ->
     -- NOTE Kind signatures already put parentheses around in all cases, so
     -- skip this layer of parentheses. The reason for this behavior is that
@@ -194,6 +177,22 @@ p_conDeclField ConDeclField {..} = do
     p_hsType (unLoc cd_fld_type)
 p_conDeclField (XConDeclField NoExt) = notImplemented "XConDeclField"
 
+tyOpTree :: LHsType GhcPs -> OpTree (LHsType GhcPs) (Located RdrName)
+tyOpTree (L _ (HsOpTy NoExt l op r)) =
+  OpBranch (tyOpTree l) op (tyOpTree r)
+tyOpTree n = OpNode n
+
+p_tyOpTree :: OpTree (LHsType GhcPs) (Located RdrName) -> R ()
+p_tyOpTree (OpNode n) = located n p_hsType
+p_tyOpTree (OpBranch l op r) = do
+  switchLayout [opTreeLoc l] $ do
+    p_tyOpTree l
+  breakpoint
+  inci . switchLayout [opTreeLoc r] $ do
+    p_rdrName op
+    space
+    p_tyOpTree r
+
 ----------------------------------------------------------------------------
 -- Conversion functions
 

src/Ormolu/Printer/Operators.hs

@@ -0,0 +1,184 @@
+{-# LANGUAGE LambdaCase          #-}
+{-# LANGUAGE ScopedTypeVariables #-}
+
+-- | This module helps handle operator chains composed of different
+-- operators that may have different precedence and fixities.
+
+module Ormolu.Printer.Operators
+  ( OpTree (..)
+  , opTreeLoc
+  , reassociateOpTree
+  ) where
+
+import BasicTypes (Fixity (..), SourceText (NoSourceText), defaultFixity, compareFixity)
+import Data.Function (on)
+import Data.List
+import Data.Maybe (fromMaybe)
+import GHC
+import OccName (mkVarOcc)
+import RdrName (mkRdrUnqual)
+import SrcLoc (combineSrcSpans)
+
+-- | Intermediate representation of operator trees. It has two type
+-- parameters: @ty@ is the type of sub-expressions, while @op@ is the type
+-- of operators.
+
+data OpTree ty op
+  = OpNode ty
+  | OpBranch
+      (OpTree ty op)
+      op
+      (OpTree ty op)
+
+-- | Return combined 'SrcSpan's of all elements in this 'OpTree'.
+
+opTreeLoc :: OpTree (Located a) b -> SrcSpan
+opTreeLoc (OpNode (L l _)) = l
+opTreeLoc (OpBranch l _ r) = combineSrcSpans (opTreeLoc l) (opTreeLoc r)
+
+-- | Re-associate an 'OpTree' taking into account automagically inferred
+-- relative precedence of operators. Users are expected to first construct
+-- an initial 'OpTree', then re-associate it using this function before
+-- printing.
+
+reassociateOpTree
+  :: (op -> Maybe RdrName)      -- ^ How to get name of an operator
+  -> OpTree (Located ty) (Located op) -- ^ Original 'OpTree'
+  -> OpTree (Located ty) (Located op) -- ^ Re-associated 'OpTree'
+reassociateOpTree getOpName opTree =
+  reassociateOpTreeWith
+    (buildFixityMap getOpName normOpTree)
+    (getOpName . unLoc)
+    normOpTree
+  where
+    normOpTree = normalizeOpTree opTree
+
+-- | Re-associate an 'OpTree' given the map with operator fixities.
+
+reassociateOpTreeWith
+  :: forall ty op.
+     [(RdrName, Fixity)]        -- ^ Fixity map for operators
+  -> (op -> Maybe RdrName)      -- ^ How to get the name of an operator
+  -> OpTree ty op               -- ^ Original 'OpTree'
+  -> OpTree ty op               -- ^ Re-associated 'OpTree'
+reassociateOpTreeWith fixityMap getOpName = go
+  where
+    fixityOf :: op -> Fixity
+    fixityOf op = fromMaybe defaultFixity $ do
+      opName <- getOpName op
+      lookup opName fixityMap
+
+    -- Here, left branch is already associated and the root alongside with
+    -- the right branch is right-associated. This function picks up one item
+    -- from the right and inserts it correctly to the left.
+    --
+    -- Also, we are using the 'compareFixity' function which returns if the
+    -- expression should associate to right.
+    go :: OpTree ty op -> OpTree ty op
+    -- base cases
+    go t@(OpNode _) = t
+    go t@(OpBranch (OpNode _) _ (OpNode _)) = t
+    -- shift one operator to the left at the beginning
+    go (OpBranch l@(OpNode _) op (OpBranch l' op' r')) =
+      go (OpBranch (OpBranch l op l') op' r')
+    -- at the last operator, place the operator and don't recurse
+    go (OpBranch (OpBranch l op r) op' r'@(OpNode _)) =
+      if snd $ compareFixity (fixityOf op) (fixityOf op')
+      then OpBranch l op (go $ OpBranch r op' r')
+      else OpBranch (OpBranch l op r) op' r'
+    -- else, shift one operator to left and recurse.
+    go (OpBranch (OpBranch l op r) op' (OpBranch l' op'' r')) =
+      if snd $ compareFixity (fixityOf op) (fixityOf op')
+      then go $ OpBranch (OpBranch l op (go $ OpBranch r op' l')) op'' r'
+      else go $ OpBranch (OpBranch (OpBranch l op r) op' l') op'' r'
+
+-- | Build a map of inferred 'Fixity's from an 'OpTree'.
+
+buildFixityMap
+  :: forall ty op.
+     (op -> Maybe RdrName)      -- ^ How to get the name of an operator
+  -> OpTree (Located ty) (Located op) -- ^ Operator tree
+  -> [(RdrName, Fixity)]               -- ^ Fixity map
+buildFixityMap getOpName opTree =
+  concatMap (\(i, ns) -> map (\(n, _) -> (n, fixity i InfixL)) ns)
+    . zip [0..]
+    . groupBy (doubleWithinEps 0.00001 `on` snd)
+    . (overrides ++)
+    . avgScores
+    $ score opTree
+  where
+    -- Add a special case for ($), since it is pretty unlikely for someone
+    -- to override it.
+    overrides :: [(RdrName, Double)]
+    overrides =
+      [ (mkRdrUnqual $ mkVarOcc "$", -1)
+      ]
+
+    -- Assign scores to operators based on their location in the source.
+    score :: OpTree (Located ty) (Located op) -> [(RdrName, Double)]
+    score (OpNode _) = []
+    score (OpBranch l o r) = fromMaybe (score r) $ do
+      -- If we fail to get any of these, 'defaultFixity' will be used by
+      -- 'reassociateOpTreeWith'.
+      le <- srcSpanEndLine <$> unSrcSpan (opTreeLoc l) -- left end
+      ob <- srcSpanStartLine <$> unSrcSpan (getLoc o) -- operator begin
+      oe <- srcSpanEndLine <$> unSrcSpan (getLoc o) -- operator end
+      rb <- srcSpanStartLine <$> unSrcSpan (opTreeLoc r) -- right begin
+      oc <- srcSpanStartCol <$> unSrcSpan (getLoc o) -- operator column
+      opName <- getOpName (unLoc o)
+
+      let s =
+           if le < ob
+             -- if the operator is in the beginning of a line, assign
+             -- a score relative to its column within range [0, 1).
+             then fromIntegral oc / fromIntegral (maxCol + 1)
+             -- if the operator is in the end of the line, assign the
+             -- score 1.
+             else
+               if oe < rb
+                 then 1
+                 else 2 -- otherwise, assign a high score.
+      return $ (opName, s) : score r
+
+    avgScores :: [(RdrName, Double)] -> [(RdrName, Double)]
+    avgScores
+      = sortOn snd
+        . map (\xs@((n, _):_) -> (n, avg $ map snd xs))
+        . groupBy ((==) `on` fst)
+        . sort
+
+    avg :: [Double] -> Double
+    avg i = sum i / fromIntegral (length i)
+
+    -- The start column of the rightmost operator.
+    maxCol = go opTree
+      where
+        go (OpNode (L _ _)) = 0
+        go (OpBranch l (L o _) r) = maximum
+          [ go l
+          , maybe 0 srcSpanStartCol (unSrcSpan o)
+          , go r
+          ]
+
+    unSrcSpan (RealSrcSpan r) = Just r
+    unSrcSpan (UnhelpfulSpan _) = Nothing
+
+----------------------------------------------------------------------------
+-- Helpers
+
+-- | Convert an 'OpTree' to with all operators having the same fixity and
+-- associativity (left infix).
+
+normalizeOpTree :: OpTree ty op -> OpTree ty op
+normalizeOpTree (OpNode n)
+  = OpNode n
+normalizeOpTree (OpBranch (OpNode l) lop r)
+  = OpBranch (OpNode l) lop (normalizeOpTree r)
+normalizeOpTree (OpBranch (OpBranch l' lop' r') lop r)
+  = normalizeOpTree (OpBranch l' lop' (OpBranch r' lop r))
+
+fixity :: Int -> FixityDirection -> Fixity
+fixity = Fixity NoSourceText
+
+doubleWithinEps :: Double -> Double -> Double -> Bool
+doubleWithinEps eps a b = abs (a - b) < eps