{-# LANGUAGE MultiWayIf #-}
{-# LANGUAGE ScopedTypeVariables #-}

-- | This module helps handle operator chains composed of different
-- operators that may have different precedence and fixities.
module Ormolu.Printer.Operators
  ( OpTree (..),
    OpInfo (..),
    opTreeLoc,
    reassociateOpTree,
    isHardSplitterOp,
  )
where

import Control.Applicative ((<|>))
import qualified Data.List.NonEmpty as NE
import qualified Data.Map.Strict as Map
import Data.Maybe (fromMaybe)
import GHC.Types.Name.Reader
import GHC.Types.SrcLoc
import Ormolu.Fixity
import Ormolu.Utils

-- | Intermediate representation of operator trees, where a branching is not
-- just a binary branching (with a left node, right node, and operator like
-- in the GHC's AST), but rather a n-ary branching, with n + 1 nodes and n
-- operators (n >= 1).
--
-- This representation allows us to put all the operators with the same
-- precedence level as direct siblings in this tree, to better represent the
-- idea of a chain of operators.
data OpTree ty op
  = -- | A node which is not an operator application
    OpNode ty
  | -- | A subtree of operator application(s); the invariant is: @length
    -- exprs == length ops + 1@. @OpBranches [x, y, z] [op1, op2]@
    -- represents the expression @x op1 y op2 z@.
    OpBranches [OpTree ty op] [op]
  deriving (OpTree ty op -> OpTree ty op -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
forall ty op.
(Eq ty, Eq op) =>
OpTree ty op -> OpTree ty op -> Bool
/= :: OpTree ty op -> OpTree ty op -> Bool
$c/= :: forall ty op.
(Eq ty, Eq op) =>
OpTree ty op -> OpTree ty op -> Bool
== :: OpTree ty op -> OpTree ty op -> Bool
$c== :: forall ty op.
(Eq ty, Eq op) =>
OpTree ty op -> OpTree ty op -> Bool
Eq, Int -> OpTree ty op -> ShowS
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
forall ty op. (Show ty, Show op) => Int -> OpTree ty op -> ShowS
forall ty op. (Show ty, Show op) => [OpTree ty op] -> ShowS
forall ty op. (Show ty, Show op) => OpTree ty op -> String
showList :: [OpTree ty op] -> ShowS
$cshowList :: forall ty op. (Show ty, Show op) => [OpTree ty op] -> ShowS
show :: OpTree ty op -> String
$cshow :: forall ty op. (Show ty, Show op) => OpTree ty op -> String
showsPrec :: Int -> OpTree ty op -> ShowS
$cshowsPrec :: forall ty op. (Show ty, Show op) => Int -> OpTree ty op -> ShowS
Show)

-- | Wrapper for an operator, carrying information about its name and
-- fixity.
data OpInfo op = OpInfo
  { -- | The actual operator
    forall op. OpInfo op -> op
opiOp :: op,
    -- | Its name, if available. We use 'Maybe OpName' here instead of 'OpName'
    -- because the name-fetching function received by 'reassociateOpTree'
    -- returns a 'Maybe'
    forall op. OpInfo op -> Maybe OpName
opiName :: Maybe OpName,
    -- | Information about the fixity direction and precedence level of the
    -- operator
    forall op. OpInfo op -> FixityInfo
opiFix :: FixityInfo
  }
  deriving (OpInfo op -> OpInfo op -> Bool
forall op. Eq op => OpInfo op -> OpInfo op -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: OpInfo op -> OpInfo op -> Bool
$c/= :: forall op. Eq op => OpInfo op -> OpInfo op -> Bool
== :: OpInfo op -> OpInfo op -> Bool
$c== :: forall op. Eq op => OpInfo op -> OpInfo op -> Bool
Eq)

-- | Compare the precedence level of two operators. 'OpInfo' is required
-- (and not just 'FixityInfo') because operator names are used in the case
-- of equality.
compareOp :: OpInfo op -> OpInfo op -> Maybe Ordering
compareOp :: forall op. OpInfo op -> OpInfo op -> Maybe Ordering
compareOp
  (OpInfo op
_ Maybe OpName
mName1 FixityInfo {fiMinPrecedence :: FixityInfo -> Int
fiMinPrecedence = Int
min1, fiMaxPrecedence :: FixityInfo -> Int
fiMaxPrecedence = Int
max1})
  (OpInfo op
_ Maybe OpName
mName2 FixityInfo {fiMinPrecedence :: FixityInfo -> Int
fiMinPrecedence = Int
min2, fiMaxPrecedence :: FixityInfo -> Int
fiMaxPrecedence = Int
max2}) =
    if
        -- Only declare two precedence levels as equal when
        --  * either both precedence levels are precise
        --    (fiMinPrecedence == fiMaxPrecedence) and match
        --  * or when the precedence levels are imprecise but when the
        --    operator names match
        | Int
min1 forall a. Eq a => a -> a -> Bool
== Int
min2
            Bool -> Bool -> Bool
&& Int
max1 forall a. Eq a => a -> a -> Bool
== Int
max2
            Bool -> Bool -> Bool
&& (Int
min1 forall a. Eq a => a -> a -> Bool
== Int
max1 Bool -> Bool -> Bool
|| Bool
sameSymbol) ->
            forall a. a -> Maybe a
Just Ordering
EQ
        | Int
max1 forall a. Ord a => a -> a -> Bool
< Int
min2 -> forall a. a -> Maybe a
Just Ordering
LT
        | Int
max2 forall a. Ord a => a -> a -> Bool
< Int
min1 -> forall a. a -> Maybe a
Just Ordering
GT
        | Bool
otherwise -> forall a. Maybe a
Nothing
    where
      sameSymbol :: Bool
sameSymbol = case (Maybe OpName
mName1, Maybe OpName
mName2) of
        (Just OpName
n1, Just OpName
n2) -> OpName
n1 forall a. Eq a => a -> a -> Bool
== OpName
n2
        (Maybe OpName, Maybe OpName)
_ -> Bool
False

-- | Return combined 'SrcSpan's of all elements in this 'OpTree'.
opTreeLoc :: (HasSrcSpan l) => OpTree (GenLocated l a) b -> SrcSpan
opTreeLoc :: forall l a b. HasSrcSpan l => OpTree (GenLocated l a) b -> SrcSpan
opTreeLoc (OpNode GenLocated l a
n) = forall l a. HasSrcSpan l => GenLocated l a -> SrcSpan
getLoc' GenLocated l a
n
opTreeLoc (OpBranches [OpTree (GenLocated l a) b]
exprs [b]
_) =
  NonEmpty SrcSpan -> SrcSpan
combineSrcSpans' forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. [a] -> NonEmpty a
NE.fromList forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall l a b. HasSrcSpan l => OpTree (GenLocated l a) b -> SrcSpan
opTreeLoc forall a b. (a -> b) -> a -> b
$ [OpTree (GenLocated l a) b]
exprs

-- | Re-associate an 'OpTree' taking into account precedence of operators.
-- Users are expected to first construct an initial 'OpTree', then
-- re-associate it using this function before printing.
reassociateOpTree ::
  -- | How to get name of an operator
  (op -> Maybe RdrName) ->
  -- | Fixity overrides
  FixityMap ->
  -- | Fixity Map
  LazyFixityMap ->
  -- | Original 'OpTree'
  OpTree ty op ->
  -- | Re-associated 'OpTree', with added context and info around operators
  OpTree ty (OpInfo op)
reassociateOpTree :: forall op ty.
(op -> Maybe RdrName)
-> FixityMap
-> LazyFixityMap
-> OpTree ty op
-> OpTree ty (OpInfo op)
reassociateOpTree op -> Maybe RdrName
getOpName FixityMap
fixityOverrides LazyFixityMap
fixityMap =
  forall ty op. OpTree ty (OpInfo op) -> OpTree ty (OpInfo op)
reassociateFlatOpTree
    forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall ty op. OpTree ty op -> OpTree ty op
makeFlatOpTree
    forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall op ty.
FixityMap
-> LazyFixityMap
-> (op -> Maybe RdrName)
-> OpTree ty op
-> OpTree ty (OpInfo op)
addFixityInfo FixityMap
fixityOverrides LazyFixityMap
fixityMap op -> Maybe RdrName
getOpName

-- | Wrap every operator of the tree with 'OpInfo' to carry the information
-- about its fixity (extracted from the specified fixity map).
addFixityInfo ::
  -- | Fixity overrides
  FixityMap ->
  -- | Fixity map for operators
  LazyFixityMap ->
  -- | How to get the name of an operator
  (op -> Maybe RdrName) ->
  -- | 'OpTree'
  OpTree ty op ->
  -- | 'OpTree', with fixity info wrapped around each operator
  OpTree ty (OpInfo op)
addFixityInfo :: forall op ty.
FixityMap
-> LazyFixityMap
-> (op -> Maybe RdrName)
-> OpTree ty op
-> OpTree ty (OpInfo op)
addFixityInfo FixityMap
_ LazyFixityMap
_ op -> Maybe RdrName
_ (OpNode ty
n) = forall ty op. ty -> OpTree ty op
OpNode ty
n
addFixityInfo FixityMap
fixityOverrides LazyFixityMap
fixityMap op -> Maybe RdrName
getOpName (OpBranches [OpTree ty op]
exprs [op]
ops) =
  forall ty op. [OpTree ty op] -> [op] -> OpTree ty op
OpBranches
    (forall op ty.
FixityMap
-> LazyFixityMap
-> (op -> Maybe RdrName)
-> OpTree ty op
-> OpTree ty (OpInfo op)
addFixityInfo FixityMap
fixityOverrides LazyFixityMap
fixityMap op -> Maybe RdrName
getOpName forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [OpTree ty op]
exprs)
    (op -> OpInfo op
toOpInfo forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [op]
ops)
  where
    toOpInfo :: op -> OpInfo op
toOpInfo op
o = forall op. op -> Maybe OpName -> FixityInfo -> OpInfo op
OpInfo op
o Maybe OpName
mName FixityInfo
fixityInfo
      where
        mName :: Maybe OpName
mName = OccName -> OpName
occOpName forall b c a. (b -> c) -> (a -> b) -> a -> c
. RdrName -> OccName
rdrNameOcc forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> op -> Maybe RdrName
getOpName op
o
        fixityInfo :: FixityInfo
fixityInfo =
          forall a. a -> Maybe a -> a
fromMaybe
            FixityInfo
defaultFixityInfo
            ( do
                OpName
name <- Maybe OpName
mName
                forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup OpName
name FixityMap
fixityOverrides forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> OpName -> LazyFixityMap -> Maybe FixityInfo
lookupFixity OpName
name LazyFixityMap
fixityMap
            )

-- | Given a 'OpTree' of any shape, produce a flat 'OpTree', where every
-- node and operator is directly connected to the root.
makeFlatOpTree :: OpTree ty op -> OpTree ty op
makeFlatOpTree :: forall ty op. OpTree ty op -> OpTree ty op
makeFlatOpTree (OpNode ty
n) = forall ty op. ty -> OpTree ty op
OpNode ty
n
makeFlatOpTree (OpBranches [OpTree ty op]
exprs [op]
ops) =
  forall ty op. [OpTree ty op] -> [op] -> OpTree ty op
OpBranches [OpTree ty op]
rExprs [op]
rOps
  where
    makeFlatOpTree' :: OpTree ty a -> ([OpTree ty a], [a])
makeFlatOpTree' OpTree ty a
expr = case forall ty op. OpTree ty op -> OpTree ty op
makeFlatOpTree OpTree ty a
expr of
      OpNode ty
n -> ([forall ty op. ty -> OpTree ty op
OpNode ty
n], [])
      OpBranches [OpTree ty a]
noptExprs [a]
noptOps -> ([OpTree ty a]
noptExprs, [a]
noptOps)
    flattenedSubTrees :: [([OpTree ty op], [op])]
flattenedSubTrees = forall {ty} {a}. OpTree ty a -> ([OpTree ty a], [a])
makeFlatOpTree' forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [OpTree ty op]
exprs
    rExprs :: [OpTree ty op]
rExprs = forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap forall a b. (a, b) -> a
fst [([OpTree ty op], [op])]
flattenedSubTrees
    rOps :: [op]
rOps = forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat forall a b. (a -> b) -> a -> b
$ forall {a}. [a] -> [a] -> [a]
interleave (forall a b. (a, b) -> b
snd forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [([OpTree ty op], [op])]
flattenedSubTrees) (forall (f :: * -> *) a. Applicative f => a -> f a
pure forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [op]
ops)
    interleave :: [a] -> [a] -> [a]
interleave (a
x : [a]
xs) (a
y : [a]
ys) = a
x forall a. a -> [a] -> [a]
: a
y forall a. a -> [a] -> [a]
: [a] -> [a] -> [a]
interleave [a]
xs [a]
ys
    interleave [] [a]
ys = [a]
ys
    interleave [a]
xs [] = [a]
xs

-- | Starting from a flat 'OpTree' (i.e. a n-ary tree of depth 1,
-- without regard for operator fixities), build an 'OpTree' with proper
-- sub-trees (according to the fixity info carried by the nodes).
--
-- We have two complementary ways to build the proper sub-trees:
--
-- * if we can find a set of operators "minOps" at the current level where
--     forall (op1, op2) \in minOps x minOps, op1 `equal` op2
--     forall (op1, op2) \in minOps x (opsOfCurrentLevel \ minOps),
--       op1 `lessThan` op2
--   then we can build a subtree with the exprs and ops located "between"
--   each element of minOps.
--   For example, if minOps = {op0, op2, op5},
--   and if [...] means "extract a subtree", then
--   currentLevel =
--     [ex0 op0 ex1 op1 ex2 op2 ex3 op3 ex4 op4 ex5 op5 ex6 op6 ex7]
--   will become
--     [ex0 op0 [ex1 op1 ex2] op2 [ex3 op3 ex4 op4 ex5] op5 [ex6 op6 ex7]]
-- * if we can find a set of operators "maxOps" at the current level where
--     forall (op1, op2) \in maxOps x maxOps, op1 `equal` op2
--     forall (op1, op2) \in maxOps x (opsOfCurrentLevel \ maxOps),
--       op1 `greaterThan` op2
--   then we can build a subtree with every contiguous range of elements
--   from maxOps (and the exprs on their sides)
--   For example, if maxOps = {op0, op1, op4},
--   and if [...] means "extract a subtree", then
--   currentLevel =
--     [ex0 op0 ex1 op1 ex2 op2 ex3 op3 ex4 op4 ex5 op5 ex6 op6 ex7]
--   will become
--     [[ex0 op0 ex1 op1 ex2] op2 ex3 op3 [ex4 op4 ex5] op5 ex6 op6 ex7]
--
-- We will also recursively apply the same logic on every sub-tree built
-- during the process. The two principles are not overlapping and thus are
-- required, because we are comparing precedence level ranges. In the case
-- where we can't find a non-empty set {min,max}Ops with one logic or the
-- other, we finally try to split the tree on “hard splitters” if there is
-- any.
reassociateFlatOpTree ::
  -- | Flat 'OpTree', with fixity info wrapped around each operator
  OpTree ty (OpInfo op) ->
  -- | Re-associated 'OpTree', with fixity info wrapped around each operator
  OpTree ty (OpInfo op)
reassociateFlatOpTree :: forall ty op. OpTree ty (OpInfo op) -> OpTree ty (OpInfo op)
reassociateFlatOpTree tree :: OpTree ty (OpInfo op)
tree@(OpNode ty
_) = OpTree ty (OpInfo op)
tree
reassociateFlatOpTree tree :: OpTree ty (OpInfo op)
tree@(OpBranches [OpTree ty (OpInfo op)]
noptExprs [OpInfo op]
noptOps) =
  case forall {op}. [OpInfo op] -> (Maybe [Int], Maybe [Int])
indexOfMinMaxPrecOps [OpInfo op]
noptOps of
    (Just [Int]
minIndices, Maybe [Int]
_) -> forall {ty} {op}.
[OpTree ty (OpInfo op)]
-> [OpInfo op] -> [Int] -> OpTree ty (OpInfo op)
splitTree [OpTree ty (OpInfo op)]
noptExprs [OpInfo op]
noptOps [Int]
minIndices
    (Maybe [Int]
_, Just [Int]
maxIndices) -> forall {ty} {op}.
[OpTree ty (OpInfo op)]
-> [OpInfo op] -> [Int] -> OpTree ty (OpInfo op)
groupTree [OpTree ty (OpInfo op)]
noptExprs [OpInfo op]
noptOps [Int]
maxIndices
    (Maybe [Int], Maybe [Int])
_ -> case [Int]
indicesOfHardSplitter of
      [] -> OpTree ty (OpInfo op)
tree
      [Int]
indices -> forall {ty} {op}.
[OpTree ty (OpInfo op)]
-> [OpInfo op] -> [Int] -> OpTree ty (OpInfo op)
splitTree [OpTree ty (OpInfo op)]
noptExprs [OpInfo op]
noptOps [Int]
indices
  where
    indicesOfHardSplitter :: [Int]
indicesOfHardSplitter =
      forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a b. (a, b) -> a
fst forall a b. (a -> b) -> a -> b
$
        forall a. (a -> Bool) -> [a] -> [a]
filter (FixityInfo -> Bool
isHardSplitterOp forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall op. OpInfo op -> FixityInfo
opiFix forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> b
snd) forall a b. (a -> b) -> a -> b
$
          forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0 ..] [OpInfo op]
noptOps
    indexOfMinMaxPrecOps :: [OpInfo op] -> (Maybe [Int], Maybe [Int])
indexOfMinMaxPrecOps [] = (forall a. Maybe a
Nothing, forall a. Maybe a
Nothing)
    indexOfMinMaxPrecOps (OpInfo op
oo : [OpInfo op]
oos) = forall op.
[OpInfo op]
-> Int
-> OpInfo op
-> Maybe [Int]
-> OpInfo op
-> Maybe [Int]
-> (Maybe [Int], Maybe [Int])
go [OpInfo op]
oos Int
1 OpInfo op
oo (forall a. a -> Maybe a
Just [Int
0]) OpInfo op
oo (forall a. a -> Maybe a
Just [Int
0])
      where
        go ::
          -- Remaining operators to look up
          [OpInfo op] ->
          -- Index of the next operator
          Int ->
          -- representative of the current minOps set, if there is one,
          -- or representative of the lowest precedence level encountered
          -- so far otherwise
          OpInfo op ->
          -- indices of the elements of the candidate minOps set,
          -- if there is any
          Maybe [Int] ->
          -- representative of the current maxOps set, if there is one, or
          -- representative of the highest precedence level encountered
          -- so far otherwise
          OpInfo op ->
          -- indices of the elements of the candidate maxOps set,
          -- if there is any
          Maybe [Int] ->
          -- (indices of minOps elements, indices of maxOps elements)
          (Maybe [Int], Maybe [Int])
        go :: forall op.
[OpInfo op]
-> Int
-> OpInfo op
-> Maybe [Int]
-> OpInfo op
-> Maybe [Int]
-> (Maybe [Int], Maybe [Int])
go [] Int
_ OpInfo op
_ Maybe [Int]
minRes OpInfo op
_ Maybe [Int]
maxRes = (forall a. [a] -> [a]
reverse forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe [Int]
minRes, forall a. [a] -> [a]
reverse forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe [Int]
maxRes)
        go (OpInfo op
o : [OpInfo op]
os) Int
i OpInfo op
minOpi Maybe [Int]
minRes OpInfo op
maxOpi Maybe [Int]
maxRes =
          let (OpInfo op
minOpi', Maybe [Int]
minRes') = case forall op. OpInfo op -> OpInfo op -> Maybe Ordering
compareOp OpInfo op
o OpInfo op
minOpi of
                Just Ordering
EQ -> (OpInfo op
minOpi, (:) Int
i forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe [Int]
minRes)
                Just Ordering
LT -> (OpInfo op
o, forall a. a -> Maybe a
Just [Int
i])
                Just Ordering
GT -> (OpInfo op
minOpi, Maybe [Int]
minRes)
                Maybe Ordering
Nothing -> (forall {op} {op}. OpInfo op -> OpInfo op -> OpInfo op
combine OpInfo op
minOpi OpInfo op
o, forall a. Maybe a
Nothing)
              (OpInfo op
maxOpi', Maybe [Int]
maxRes') = case forall op. OpInfo op -> OpInfo op -> Maybe Ordering
compareOp OpInfo op
o OpInfo op
maxOpi of
                Just Ordering
EQ -> (OpInfo op
maxOpi, (:) Int
i forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe [Int]
maxRes)
                Just Ordering
LT -> (OpInfo op
maxOpi, Maybe [Int]
maxRes)
                Just Ordering
GT -> (OpInfo op
o, forall a. a -> Maybe a
Just [Int
i])
                Maybe Ordering
Nothing -> (forall {op} {op}. OpInfo op -> OpInfo op -> OpInfo op
combine OpInfo op
maxOpi OpInfo op
o, forall a. Maybe a
Nothing)
              -- Merge two potential {min/max}Ops representatives for
              -- which the comparison gave 'OpUnknown' into a representative
              -- of the {lowest/highest} precedence level encountered so far
              combine :: OpInfo op -> OpInfo op -> OpInfo op
combine (OpInfo op
x Maybe OpName
_ FixityInfo
fix1) (OpInfo op
_ Maybe OpName
_ FixityInfo
fix2) =
                forall op. op -> Maybe OpName -> FixityInfo -> OpInfo op
OpInfo op
x forall a. Maybe a
Nothing (FixityInfo
fix1 forall a. Semigroup a => a -> a -> a
<> FixityInfo
fix2)
           in forall op.
[OpInfo op]
-> Int
-> OpInfo op
-> Maybe [Int]
-> OpInfo op
-> Maybe [Int]
-> (Maybe [Int], Maybe [Int])
go [OpInfo op]
os (Int
i forall a. Num a => a -> a -> a
+ Int
1) OpInfo op
minOpi' Maybe [Int]
minRes' OpInfo op
maxOpi' Maybe [Int]
maxRes'
    -- If indices = [0, 2, 5], transform
    --   [ex0 op0 ex1 op1 ex2 op2 ex3 op3 ex4 op4 ex5 op5 ex6 op6 ex7]
    -- into
    --   [ex0 op0 [ex1 op1 ex2] op2 [ex3 op3 ex4 op4 ex5] op5 [ex6 op6 ex7]]
    splitTree :: [OpTree ty (OpInfo op)]
-> [OpInfo op] -> [Int] -> OpTree ty (OpInfo op)
splitTree [OpTree ty (OpInfo op)]
nExprs [OpInfo op]
nOps [Int]
indices = forall ty op.
[OpTree ty (OpInfo op)]
-> [OpInfo op]
-> [Int]
-> Int
-> [OpTree ty (OpInfo op)]
-> [OpInfo op]
-> [OpTree ty (OpInfo op)]
-> [OpInfo op]
-> OpTree ty (OpInfo op)
go [OpTree ty (OpInfo op)]
nExprs [OpInfo op]
nOps [Int]
indices Int
0 [] [] [] []
      where
        go ::
          -- Remaining exprs to look up
          [OpTree ty (OpInfo op)] ->
          -- Remaining ops to look up
          [OpInfo op] ->
          -- Remaining list of indices of operators on which to split
          -- (sorted)
          [Int] ->
          -- Index of the next expr/op
          Int ->
          -- Bag for exprs for the subtree we are building
          [OpTree ty (OpInfo op)] ->
          -- Bag for ops for the subtree we are building
          [OpInfo op] ->
          -- Bag for exprs of the result tree
          [OpTree ty (OpInfo op)] ->
          -- Bag for ops of the result tree
          [OpInfo op] ->
          -- Result tree
          OpTree ty (OpInfo op)
        go :: forall ty op.
[OpTree ty (OpInfo op)]
-> [OpInfo op]
-> [Int]
-> Int
-> [OpTree ty (OpInfo op)]
-> [OpInfo op]
-> [OpTree ty (OpInfo op)]
-> [OpInfo op]
-> OpTree ty (OpInfo op)
go [] [OpInfo op]
_ [Int]
_ Int
_ [OpTree ty (OpInfo op)]
subExprs [OpInfo op]
subOps [OpTree ty (OpInfo op)]
resExprs [OpInfo op]
resOps =
          -- No expr left to process.
          -- because we are in a "splitting" logic, there is at least one
          -- expr in the subExprs bag, so we build a subtree (if necessary)
          -- with sub-bags, add the node/subtree to the result bag, and then
          -- emit the result tree
          let resExpr :: OpTree ty (OpInfo op)
resExpr = forall {ty} {op}.
[OpTree ty (OpInfo op)] -> [OpInfo op] -> OpTree ty (OpInfo op)
buildFromSub [OpTree ty (OpInfo op)]
subExprs [OpInfo op]
subOps
           in forall ty op. [OpTree ty op] -> [op] -> OpTree ty op
OpBranches (forall a. [a] -> [a]
reverse (OpTree ty (OpInfo op)
resExpr forall a. a -> [a] -> [a]
: [OpTree ty (OpInfo op)]
resExprs)) (forall a. [a] -> [a]
reverse [OpInfo op]
resOps)
        go (OpTree ty (OpInfo op)
x : [OpTree ty (OpInfo op)]
xs) (OpInfo op
o : [OpInfo op]
os) (Int
idx : [Int]
idxs) Int
i [OpTree ty (OpInfo op)]
subExprs [OpInfo op]
subOps [OpTree ty (OpInfo op)]
resExprs [OpInfo op]
resOps
          | Int
i forall a. Eq a => a -> a -> Bool
== Int
idx =
              -- The op we are looking at is one on which we need to split.
              -- So we build a subtree from the sub-bags and the current
              -- expr, append it to the result exprs, and continue with
              -- cleared sub-bags
              let resExpr :: OpTree ty (OpInfo op)
resExpr = forall {ty} {op}.
[OpTree ty (OpInfo op)] -> [OpInfo op] -> OpTree ty (OpInfo op)
buildFromSub (OpTree ty (OpInfo op)
x forall a. a -> [a] -> [a]
: [OpTree ty (OpInfo op)]
subExprs) [OpInfo op]
subOps
               in forall ty op.
[OpTree ty (OpInfo op)]
-> [OpInfo op]
-> [Int]
-> Int
-> [OpTree ty (OpInfo op)]
-> [OpInfo op]
-> [OpTree ty (OpInfo op)]
-> [OpInfo op]
-> OpTree ty (OpInfo op)
go [OpTree ty (OpInfo op)]
xs [OpInfo op]
os [Int]
idxs (Int
i forall a. Num a => a -> a -> a
+ Int
1) [] [] (OpTree ty (OpInfo op)
resExpr forall a. a -> [a] -> [a]
: [OpTree ty (OpInfo op)]
resExprs) (OpInfo op
o forall a. a -> [a] -> [a]
: [OpInfo op]
resOps)
        go (OpTree ty (OpInfo op)
x : [OpTree ty (OpInfo op)]
xs) [OpInfo op]
ops [Int]
idxs Int
i [OpTree ty (OpInfo op)]
subExprs [OpInfo op]
subOps [OpTree ty (OpInfo op)]
resExprs [OpInfo op]
resOps =
          -- Either there is no op left, or the op we are looking at is not
          -- one on which we need to split. So we just add both the current
          -- expr and current op (if there is any) to the sub-bags
          let ([OpInfo op]
ops', [OpInfo op]
subOps') = forall {a}. [a] -> [a] -> ([a], [a])
moveOneIfPossible [OpInfo op]
ops [OpInfo op]
subOps
           in forall ty op.
[OpTree ty (OpInfo op)]
-> [OpInfo op]
-> [Int]
-> Int
-> [OpTree ty (OpInfo op)]
-> [OpInfo op]
-> [OpTree ty (OpInfo op)]
-> [OpInfo op]
-> OpTree ty (OpInfo op)
go [OpTree ty (OpInfo op)]
xs [OpInfo op]
ops' [Int]
idxs (Int
i forall a. Num a => a -> a -> a
+ Int
1) (OpTree ty (OpInfo op)
x forall a. a -> [a] -> [a]
: [OpTree ty (OpInfo op)]
subExprs) [OpInfo op]
subOps' [OpTree ty (OpInfo op)]
resExprs [OpInfo op]
resOps

    -- If indices = [0, 1, 4], transform
    --   [ex0 op0 ex1 op1 ex2 op2 ex3 op3 ex4 op4 ex5 op5 ex6 op6 ex7]
    -- into
    --   [[ex0 op0 ex1 op1 ex2] op2 ex3 op3 [ex4 op4 ex5] op5 ex6 op6 ex7]
    groupTree :: [OpTree ty (OpInfo op)]
-> [OpInfo op] -> [Int] -> OpTree ty (OpInfo op)
groupTree [OpTree ty (OpInfo op)]
nExprs [OpInfo op]
nOps [Int]
indices = forall ty op.
[OpTree ty (OpInfo op)]
-> [OpInfo op]
-> [Int]
-> Int
-> [OpTree ty (OpInfo op)]
-> [OpInfo op]
-> [OpTree ty (OpInfo op)]
-> [OpInfo op]
-> OpTree ty (OpInfo op)
go [OpTree ty (OpInfo op)]
nExprs [OpInfo op]
nOps [Int]
indices Int
0 [] [] [] []
      where
        go ::
          -- remaining exprs to look up
          [OpTree ty (OpInfo op)] ->
          -- remaining ops to look up
          [OpInfo op] ->
          -- remaining list of indices of operators on which to group
          -- (sorted)
          [Int] ->
          -- index of the next expr/op
          Int ->
          -- bag for exprs for the subtree we are building
          [OpTree ty (OpInfo op)] ->
          -- bag for ops for the subtree we are building
          [OpInfo op] ->
          -- bag for exprs of the result tree
          [OpTree ty (OpInfo op)] ->
          -- bag for ops of the result tree
          [OpInfo op] ->
          -- result tree
          OpTree ty (OpInfo op)
        go :: forall ty op.
[OpTree ty (OpInfo op)]
-> [OpInfo op]
-> [Int]
-> Int
-> [OpTree ty (OpInfo op)]
-> [OpInfo op]
-> [OpTree ty (OpInfo op)]
-> [OpInfo op]
-> OpTree ty (OpInfo op)
go [] [OpInfo op]
_ [Int]
_ Int
_ [OpTree ty (OpInfo op)]
subExprs [OpInfo op]
subOps [OpTree ty (OpInfo op)]
resExprs [OpInfo op]
resOps =
          -- no expr left to process
          -- because we are in a "grouping" logic, the subExprs bag might be
          -- empty. If it is not, we build a subtree (if necessary) with
          -- sub-bags and add the resulting node/subtree to the result bag.
          -- In any case, we then emit the result tree
          let resExprs' :: [OpTree ty (OpInfo op)]
resExprs' =
                if forall (t :: * -> *) a. Foldable t => t a -> Bool
null [OpTree ty (OpInfo op)]
subExprs
                  then [OpTree ty (OpInfo op)]
resExprs
                  else forall {ty} {op}.
[OpTree ty (OpInfo op)] -> [OpInfo op] -> OpTree ty (OpInfo op)
buildFromSub [OpTree ty (OpInfo op)]
subExprs [OpInfo op]
subOps forall a. a -> [a] -> [a]
: [OpTree ty (OpInfo op)]
resExprs
           in forall ty op. [OpTree ty op] -> [op] -> OpTree ty op
OpBranches (forall a. [a] -> [a]
reverse [OpTree ty (OpInfo op)]
resExprs') (forall a. [a] -> [a]
reverse [OpInfo op]
resOps)
        go (OpTree ty (OpInfo op)
x : [OpTree ty (OpInfo op)]
xs) (OpInfo op
o : [OpInfo op]
os) (Int
idx : [Int]
idxs) Int
i [OpTree ty (OpInfo op)]
subExprs [OpInfo op]
subOps [OpTree ty (OpInfo op)]
resExprs [OpInfo op]
resOps
          | Int
i forall a. Eq a => a -> a -> Bool
== Int
idx =
              -- The op we are looking at is one on which we need to group.
              -- So we just add the current expr and op to the sub-bags.
              forall ty op.
[OpTree ty (OpInfo op)]
-> [OpInfo op]
-> [Int]
-> Int
-> [OpTree ty (OpInfo op)]
-> [OpInfo op]
-> [OpTree ty (OpInfo op)]
-> [OpInfo op]
-> OpTree ty (OpInfo op)
go [OpTree ty (OpInfo op)]
xs [OpInfo op]
os [Int]
idxs (Int
i forall a. Num a => a -> a -> a
+ Int
1) (OpTree ty (OpInfo op)
x forall a. a -> [a] -> [a]
: [OpTree ty (OpInfo op)]
subExprs) (OpInfo op
o forall a. a -> [a] -> [a]
: [OpInfo op]
subOps) [OpTree ty (OpInfo op)]
resExprs [OpInfo op]
resOps
        go (OpTree ty (OpInfo op)
x : [OpTree ty (OpInfo op)]
xs) [OpInfo op]
ops [Int]
idxs Int
i subExprs :: [OpTree ty (OpInfo op)]
subExprs@(OpTree ty (OpInfo op)
_ : [OpTree ty (OpInfo op)]
_) [OpInfo op]
subOps [OpTree ty (OpInfo op)]
resExprs [OpInfo op]
resOps =
          -- Either there is no op left, or the op we are looking at is not
          -- one on which we need to split, but in any case the sub-bags are
          -- not empty. So we finalize the started group using sub-bags and
          -- the current expr, to form a subtree which is then added to the
          -- result bag.
          let ([OpInfo op]
ops', [OpInfo op]
resOps') = forall {a}. [a] -> [a] -> ([a], [a])
moveOneIfPossible [OpInfo op]
ops [OpInfo op]
resOps
              resExpr :: OpTree ty (OpInfo op)
resExpr = forall {ty} {op}.
[OpTree ty (OpInfo op)] -> [OpInfo op] -> OpTree ty (OpInfo op)
buildFromSub (OpTree ty (OpInfo op)
x forall a. a -> [a] -> [a]
: [OpTree ty (OpInfo op)]
subExprs) [OpInfo op]
subOps
           in forall ty op.
[OpTree ty (OpInfo op)]
-> [OpInfo op]
-> [Int]
-> Int
-> [OpTree ty (OpInfo op)]
-> [OpInfo op]
-> [OpTree ty (OpInfo op)]
-> [OpInfo op]
-> OpTree ty (OpInfo op)
go [OpTree ty (OpInfo op)]
xs [OpInfo op]
ops' [Int]
idxs (Int
i forall a. Num a => a -> a -> a
+ Int
1) [] [] (OpTree ty (OpInfo op)
resExpr forall a. a -> [a] -> [a]
: [OpTree ty (OpInfo op)]
resExprs) [OpInfo op]
resOps'
        go (OpTree ty (OpInfo op)
x : [OpTree ty (OpInfo op)]
xs) [OpInfo op]
ops [Int]
idxs Int
i [] [OpInfo op]
subOps [OpTree ty (OpInfo op)]
resExprs [OpInfo op]
resOps =
          -- Either there is no op left, or the op we are looking at is not
          -- one on which we need to split, but the sub-bags are empty. So
          -- we just add both the current expr and current op (if there is
          -- any) to the result bags
          let ([OpInfo op]
ops', [OpInfo op]
resOps') = forall {a}. [a] -> [a] -> ([a], [a])
moveOneIfPossible [OpInfo op]
ops [OpInfo op]
resOps
           in forall ty op.
[OpTree ty (OpInfo op)]
-> [OpInfo op]
-> [Int]
-> Int
-> [OpTree ty (OpInfo op)]
-> [OpInfo op]
-> [OpTree ty (OpInfo op)]
-> [OpInfo op]
-> OpTree ty (OpInfo op)
go [OpTree ty (OpInfo op)]
xs [OpInfo op]
ops' [Int]
idxs (Int
i forall a. Num a => a -> a -> a
+ Int
1) [] [OpInfo op]
subOps (OpTree ty (OpInfo op)
x forall a. a -> [a] -> [a]
: [OpTree ty (OpInfo op)]
resExprs) [OpInfo op]
resOps'

    moveOneIfPossible :: [a] -> [a] -> ([a], [a])
moveOneIfPossible [] [a]
bs = ([], [a]
bs)
    moveOneIfPossible (a
a : [a]
as) [a]
bs = ([a]
as, a
a forall a. a -> [a] -> [a]
: [a]
bs)

    buildFromSub :: [OpTree ty (OpInfo op)] -> [OpInfo op] -> OpTree ty (OpInfo op)
buildFromSub [OpTree ty (OpInfo op)]
subExprs [OpInfo op]
subOps = forall ty op. OpTree ty (OpInfo op) -> OpTree ty (OpInfo op)
reassociateFlatOpTree forall a b. (a -> b) -> a -> b
$ case [OpTree ty (OpInfo op)]
subExprs of
      -- Do not build a subtree when the potential subtree would have
      -- 1 expr(s) and 0 op(s)
      [OpTree ty (OpInfo op)
x] -> OpTree ty (OpInfo op)
x
      [OpTree ty (OpInfo op)]
_ -> forall ty op. [OpTree ty op] -> [op] -> OpTree ty op
OpBranches (forall a. [a] -> [a]
reverse [OpTree ty (OpInfo op)]
subExprs) (forall a. [a] -> [a]
reverse [OpInfo op]
subOps)

-- | Indicate if an operator has @'InfixR' 0@ fixity. We special-case this
-- class of operators because they often have, like ('$'), a specific
-- “separator” use-case, and we sometimes format them differently than other
-- operators.
isHardSplitterOp :: FixityInfo -> Bool
isHardSplitterOp :: FixityInfo -> Bool
isHardSplitterOp = (forall a. Eq a => a -> a -> Bool
== Maybe FixityDirection -> Int -> Int -> FixityInfo
FixityInfo (forall a. a -> Maybe a
Just FixityDirection
InfixR) Int
0 Int
0)