OpTree.hs

{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE LambdaCase #-}

-- | Printing of operator trees.
module Ormolu.Printer.Meat.Declaration.OpTree
  ( p_exprOpTree,
    exprOpTree,
    p_cmdOpTree,
    cmdOpTree,
    p_tyOpTree,
    tyOpTree,
    getOpName,
    getOpNameStr,
  )
where

import Data.Functor ((<&>))
import Data.List.NonEmpty (NonEmpty (..))
import Data.List.NonEmpty qualified as NE
import GHC.Hs
import GHC.Types.Fixity
import GHC.Types.Name (occNameString)
import GHC.Types.Name.Reader (RdrName, rdrNameOcc)
import GHC.Types.SrcLoc
import Ormolu.Config (poIndentation, poTrailingSectionOperators)
import Ormolu.Parser.CommentStream (LComment)
import Ormolu.Printer.Combinators
import Ormolu.Printer.Meat.Common (p_rdrName)
import Ormolu.Printer.Meat.Declaration.Value
  ( IsApplicand (..),
    cmdTopPlacement,
    exprPlacement,
    p_hsCmdTop,
    p_hsExpr,
    p_hsExpr',
  )
import Ormolu.Printer.Meat.Type (p_hsType)
import Ormolu.Printer.Operators

-- | Extract the operator name of the specified 'HsExpr' if this expression
-- corresponds to an operator.
getOpName :: HsExpr GhcPs -> Maybe RdrName
getOpName = \case
  HsVar _ (L _ a) -> Just a
  _ -> Nothing

-- | Convert an operator name to a 'String'.
getOpNameStr :: RdrName -> String
getOpNameStr = occNameString . rdrNameOcc

-- | Decide whether the operands of an operator chain should be hanging.
opBranchPlacement ::
  (HasLoc l) =>
  -- | Placer function for nodes
  (ty -> Placement) ->
  -- | First expression of the chain
  OpTree (GenLocated l ty) op ->
  -- | Last expression of the chain
  OpTree (GenLocated l ty) op ->
  Placement
opBranchPlacement placer firstExpr lastExpr
  -- If the start of the first argument and the start of the last argument
  -- are on the same line, and the last argument has a hanging form, use
  -- hanging placement.
  | isOneLineSpan
      ( mkSrcSpan
          (srcSpanStart (opTreeLoc firstExpr))
          (srcSpanStart (opTreeLoc lastExpr))
      ),
    OpNode (L _ n) <- lastExpr =
      placer n
  | otherwise = Normal

-- | Decide whether or not to use braces based on the layout and placement
-- of an expression in an infix operator application.
opBranchBraceStyle :: Placement -> R (R () -> R ())
opBranchBraceStyle placement =
  getLayout <&> \case
    SingleLine -> useBraces
    MultiLine -> case placement of
      Hanging -> useBraces
      Normal -> dontUseBraces

-- | Convert a 'LHsExpr' containing an operator tree to the 'OpTree'
-- intermediate representation.
exprOpTree :: LHsExpr GhcPs -> OpTree (LHsExpr GhcPs) (LHsExpr GhcPs)
exprOpTree (L _ (OpApp _ x op y)) = BinaryOpBranches (exprOpTree x) op (exprOpTree y)
exprOpTree n = OpNode n

-- | Print an operator tree where leaves are values.
p_exprOpTree ::
  -- | Bracket style to use
  BracketStyle ->
  -- | N-ary 'OpTree' to render, enhanced with information regarding
  -- operator fixity
  OpTree (LHsExpr GhcPs) (OpInfo (LHsExpr GhcPs)) ->
  R ()
p_exprOpTree s (OpNode x) = located x (p_hsExpr' NotApplicand s)
p_exprOpTree s t@(OpBranches exprs@(firstExpr :| otherExprs) ops) = do
  trailingSectionOperators <- getPrinterOpt poTrailingSectionOperators

  let placement =
        opBranchPlacement
          exprPlacement
          firstExpr
          (last otherExprs)
      rightMostNode = \case
        n@(OpNode _) -> n
        OpBranches exprs'' _ -> rightMostNode (NE.last exprs'')
      isDoBlock = \case
        OpNode (L _ (HsDo _ ctx _)) -> case ctx of
          DoExpr _ -> True
          MDoExpr _ -> True
          _ -> False
        _ -> False
      -- Whether we could place the operator in a trailing position,
      -- followed by a breakpoint before the RHS
      couldBeTrailing (prevExpr, opi) =
        -- Enabled by config
        trailingSectionOperators
          -- An operator with fixity InfixR 0, like seq, $, and the $ variants,
          -- is required.
          && isHardSplitterOp (opiFixityApproximation opi)
          -- the LHS must be single-line.
          && isOneLineSpan (opTreeLoc prevExpr)
          -- This can only happen when a breakpoint would have been added
          -- anyway.
          && placement == Normal
          -- If the node just to the left of the operator (that is, the
          -- rightmost node of the subtree prevExpr) is a do-block, then we
          -- cannot place the operator in a trailing position, because it
          -- would be read as being part of the do-block.
          && not (isDoBlock $ rightMostNode prevExpr)
      -- A staircase of two or more trailing operators is only worthwhile when
      -- the operand at the very end of the chain has a hanging form (a do
      -- block, lambda, case, etc.): the trailing operator then introduces that
      -- block. When such a chain ends in an ordinary expression (a variable,
      -- literal, or plain application) the trailing layout only produces
      -- ever-deepening indentation, so we fall back to the leading-operator
      -- layout. A single hard splitter is exempt: it does not form a pyramid
      -- and trailing is the idiomatic way to introduce its operand.
      chainEndsInHangingForm =
        case rightMostNode t of
          OpNode (L _ n) -> exprPlacement n == Hanging
          _ -> False
      isSingleOperator = case ops of
        [_] -> True
        _ -> False
  -- A comment written on its own line in front of an operator forces the
  -- operator onto a line of its own. In trailing position that line starts
  -- at the indentation of the statement, and @$@ at the start of a line in
  -- a @do@ block is read as a new statement rather than as a continuation.
  -- The leading layout indents instead, so it keeps the meaning.
  opsAreCommented <-
    or <$> traverse (fmap (not . null) . leadingComments . opiOp) ops
  -- If all operators at the current level match the conditions to be
  -- trailing, and the chain is either a single operator or ends in a
  -- hanging form, then put the operators in a trailing position.
  let isTrailing =
        (isSingleOperator || chainEndsInHangingForm)
          && not opsAreCommented
          && all couldBeTrailing (zip (NE.toList exprs) ops)
  ub <- if isTrailing then return useBraces else opBranchBraceStyle placement
  indent <- getPrinterOpt poIndentation
  let p_x = ub $ p_exprOpTree s firstExpr
      putOpsExprs prevExpr (opi : ops') (expr : exprs') = do
        let isLast = null exprs'
            ub' = if not isLast then ub else id
            p_op = located (opiOp opi) p_hsExpr
            p_y = ub' $ p_exprOpTree N expr
        if isTrailing
          then do
            space
            p_op
            placeHanging
              -- When we have a chain of trailing operators (staircase style),
              -- the last operand, when multiline, is allowed to hang
              -- (ex: do block, lambda...)
              ( if isLast && (not . isOneLineSpan . opTreeLoc $ expr)
                  then opBranchPlacement exprPlacement prevExpr expr
                  else Normal
              )
              $ do
                p_y
                putOpsExprs expr ops' exprs'
          else do
            let withPlacement m =
                  -- This case prevents an operator from being indented past the start of a `do` block
                  -- constituting its left operand, thus altering the AST.
                  -- This is only relevant when the `do` block is on one line, as otherwise we will
                  -- insert a newline after `do` anyway.
                  -- This isn't an issue in Ormolu because this problem doesn't come up with 2-space
                  -- indents, only when the indentation goes past the "do" column.
                  if indent > 2 && isDoBlock firstExpr && isOneLineSpan (opTreeLoc firstExpr)
                    then breakpoint >> m
                    else placeHanging placement m
            withPlacement $ do
              p_op
              space
              p_y
            putOpsExprs expr ops' exprs'
      putOpsExprs _ _ _ = pure ()
  switchLayout [opTreeLoc t] $ do
    p_x
    putOpsExprs firstExpr ops otherExprs

-- | The comments that will be printed in front of a located thing.
leadingComments :: (HasLoc l) => GenLocated l a -> R [LComment]
leadingComments (L l _) = case locA l of
  RealSrcSpan spn _ -> getCommentsBefore spn
  _ -> pure []

-- | Convert a 'LHsCmdTop' containing an operator tree to the 'OpTree'
-- intermediate representation.
cmdOpTree :: LHsCmdTop GhcPs -> OpTree (LHsCmdTop GhcPs) (LHsExpr GhcPs)
cmdOpTree = \case
  (L _ (HsCmdTop _ (L _ (HsCmdArrForm _ op Infix [x, y])))) ->
    BinaryOpBranches (cmdOpTree x) op (cmdOpTree y)
  n -> OpNode n

-- | Print an operator tree where leaves are commands.
p_cmdOpTree ::
  -- | Bracket style to use
  BracketStyle ->
  -- | N-ary OpTree to render, enhanced with information regarding operator
  -- fixity
  OpTree (LHsCmdTop GhcPs) (OpInfo (LHsExpr GhcPs)) ->
  R ()
p_cmdOpTree s (OpNode x) = located x (p_hsCmdTop s)
p_cmdOpTree s t@(OpBranches (firstExpr :| otherExprs) ops) = do
  let placement =
        opBranchPlacement
          cmdTopPlacement
          firstExpr
          (last otherExprs)
  ub <- opBranchBraceStyle placement
  let p_x = ub $ p_cmdOpTree s firstExpr
      putOpsExprs (opi : ops') (expr : exprs') = do
        let ub' = if not (null exprs') then ub else id
            p_op = located (opiOp opi) p_hsExpr
            p_y = ub' $ p_cmdOpTree N expr
        placeHanging placement $ do
          p_op
          space
          p_y
        putOpsExprs ops' exprs'
      putOpsExprs _ _ = pure ()
  switchLayout [opTreeLoc t] $ do
    p_x
    putOpsExprs ops otherExprs

-- | Check whether the given expression has a hanging form. Added for
-- symmetry with 'exprPlacement' and 'cmdTopPlacement', all of which are used
-- in the @p_xxxOpTree@ functions together with 'opBranchPlacement'.
tyOpPlacement :: HsType GhcPs -> Placement
tyOpPlacement = \case
  _ -> Normal

-- | Convert an 'LHsType' containing an operator tree to the 'OpTree'
-- intermediate representation.
tyOpTree :: LHsType GhcPs -> OpTree (LHsType GhcPs) (LocatedN RdrName)
tyOpTree (L _ (HsOpTy _ _ l op r)) =
  BinaryOpBranches (tyOpTree l) op (tyOpTree r)
tyOpTree n = OpNode n

-- | Print an operator tree where leaves are types.
p_tyOpTree ::
  -- | N-ary 'OpTree' to render, enhanced with information regarding
  -- operator fixity
  OpTree (LHsType GhcPs) (OpInfo (LocatedN RdrName)) ->
  R ()
p_tyOpTree (OpNode n) = located n p_hsType
p_tyOpTree t@(OpBranches (firstExpr :| otherExprs) ops) = do
  let placement =
        opBranchPlacement
          tyOpPlacement
          firstExpr
          (last otherExprs)
      p_x = p_tyOpTree firstExpr
      putOpsExprs (opi : ops') (expr : exprs') = do
        let p_op = p_rdrName (opiOp opi)
            p_y = p_tyOpTree expr
        placeHanging
          placement
          $ do
            p_op
            space
            p_y
        putOpsExprs ops' exprs'
      putOpsExprs _ _ = pure ()
  switchLayout [opTreeLoc t] $ do
    ub <- opBranchBraceStyle placement
    ub p_x
    putOpsExprs ops otherExprs