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.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 if 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 beginning of the first argument and the last argument starts on
  -- the same line, and the second 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 to use braces or not 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 $ variants,
          -- is required
          && isHardSplitterOp (opiFixityApproximation opi)
          -- the LHS must be single-line
          && isOneLineSpan (opTreeLoc prevExpr)
          -- can only happen when a breakpoint would have been added anyway
          && placement == Normal
          -- if the node just on the left of the operator (so 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)
      -- If all operators at the current level match the conditions to be
      -- trailing, then put them in a trailing position
      isTrailing = 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

-- | 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 if given expression has a hanging form. Added for symmetry with
-- exprPlacement and cmdTopPlacement, which are all used in p_xxxOpTree
-- functions with opBranchPlacement.
tyOpPlacement :: HsType GhcPs -> Placement
tyOpPlacement = \case
  _ -> Normal

-- | Convert a 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