tilia-0.1.2.0: src/Tilia/Render/Expression.hs

{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE OverloadedLabels #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}

-- | Expressions, and the equations and blocks built out of them.
module Tilia.Render.Expression
  ( hsExpr,
    hsExprIn,
    valDecl,
    untypedSplice,
  )
where

import Control.Applicative ((<|>))
import Data.Choice (Choice, fromBool, isTrue, pattern Is, pattern Isn't)
import Data.Function (on)
import Data.Generics.Schemes (listify)
import Data.List (sortBy, unsnoc)
import Data.List.NonEmpty (NonEmpty (..))
import Data.List.NonEmpty qualified as NE
import Data.Map.Strict qualified as Map
import Data.Text (Text)
import Data.Text qualified as T
import GHC.Data.FastString (unpackFS)
import GHC.Hs hiding (Fixity)
import GHC.LanguageExtensions.Type (Extension (..))
import GHC.Types.Basic (Boxity (..))
import GHC.Types.Fixity (LexicalFixity (..))
import GHC.Types.Name.Occurrence (isVarOcc, occNameString)
import GHC.Types.Name.Reader (RdrName (..), mkVarUnqual, rdrNameOcc)
import GHC.Types.SourceText
import GHC.Types.SrcLoc
  ( GenLocated (..),
    isZeroWidthSpan,
    leftmost_smallest,
    noSrcSpan,
    unLoc,
  )
import Language.Haskell.Syntax.Basic (field_label)
import Tilia.Doc.Body
import Tilia.Doc.Combinators
import Tilia.Fixity
  ( Fixity,
    Namespace (..),
  )
import Tilia.Render.Context
import Tilia.Render.Layout
import Tilia.Render.Literal (stringLiteral)
import Tilia.Render.Name
import Tilia.Render.Operator
import Tilia.Render.Pattern
import Tilia.Render.Type
import Tilia.Span
import Tilia.Span.Ghc

-- | An expression.
hsExpr :: Ctx -> LHsExpr GhcPs -> Doc
hsExpr ctx = hsExprIn ctx plainSite

-- | An expression that knows where it stands.
hsExprIn :: Ctx -> Site -> LHsExpr GhcPs -> Doc
hsExprIn ctx site l = at ctx l (exprBody ctx site (spanOf l))

-- | Every shape an expression can take.
exprBody :: Ctx -> Site -> Maybe Span -> HsExpr GhcPs -> Doc
exprBody ctx site here = \case
  HsVar _ n -> name ctx n
  HsOverLabel src _ -> txt "#" <> sourceText src
  HsIPVar _ (HsIPName n) -> txt "?" <> outputable n
  HsOverLit _ v -> outputable (ol_val v)
  HsLit _ lit -> case lit of
    HsString (SourceText s) _ -> stringLiteral s
    HsStringPrim (SourceText s) _ -> stringLiteral s
    HsMultilineString (SourceText s) _ -> stringLiteral s
    other -> outputable other
  HsLam _ variant' mg -> lambda ctx site variant' (ExprBody ctx) mg
  HsApp _ f x -> application ctx site f x
  HsAppType at' e a ->
    hsExpr ctx e
      <> breakOrSpace
      <> indent
        ( atSpan
            ctx
            (spanOf at' <> spanOf (hswc_body a))
            (txt "@" <> hsTypeBody ctx (spanOf (hswc_body a)) (unLoc (hswc_body a)))
        )
  OpApp _ x op y -> exprChain ctx site x op y
  NegApp _ e _ ->
    txt "-"
      <> includeWhen (isLiteral e) (negativeLiteralGap ctx)
      <> hsExpr ctx e
  HsPar _ e ->
    layoutWithin ctx here (spanOf e) $
      parensWith (closingFor site) (insideBrackets here (apart (hsExpr ctx e)))
    where
      apart d = case unLoc e of
        SectionR _ (L _ (HsVar _ (L _ (Unqual o)))) _
          | unboxedParens, take 1 (occNameString o) == "#" -> space <> d
        SectionL _ _ (L _ (HsVar _ (L _ (Unqual o))))
          | unboxedParens, occNameString o == "#" -> d <> space
        _ -> d
      unboxedParens =
        extensionOn ctx UnboxedTuples || extensionOn ctx UnboxedSums
  SectionL _ x op -> hsExpr ctx x <> breakOrSpace <> indent (hsExpr ctx op)
  SectionR _ op x -> hsExpr ctx op <> breakOrSpace <> indent (hsExpr ctx x)
  ExplicitTuple _ args boxity -> tuple ctx here (closingFor site) boxity args
  ExplicitSum _ tag arity e -> unboxedSum (closingFor site) tag arity (hsExpr ctx e)
  HsCase anns e mg -> caseOf ctx site (ExprBody ctx) anns e mg
  HsIf anns c t e -> ifThenElse ctx (bodyIn (ExprBody ctx) (siteBracing site)) anns c t e
  HsMultiIf _ guards ->
    txt "if"
      <> breakOrSpace
      <> underSite site (sepBy breakOrSpace (fmap alternative (NE.toList guards)))
    where
      alternative =
        guardedRhs ctx Normal (siteBracing site) (ExprBody ctx) RightArrow
  HsLet tokens binds e ->
    letIn ctx (bodyIn (ExprBody ctx) (siteBracing site)) nested tokens binds e
    where
      nested = case unLoc e of
        HsLet{} -> True
        _ -> False
  HsDo anns flavour es -> case flavour of
    DoExpr moduleName -> doBlock moduleName "do"
    MDoExpr moduleName -> doBlock moduleName "mdo"
    ListComp -> comprehension ctx site here es
    MonadComp -> comprehension ctx site here es
    GhciStmtCtxt -> error "Tilia: GhciStmtCtxt cannot occur in a source file"
    where
      doBlock moduleName word =
        foldMap (\m -> outputable m <> txt ".") moduleName
          <> keywordAt ctx (doKeywordSpan anns) word
          <> statements ctx site (ExprBody ctx) es
  ExplicitList _ xs ->
    bracketsWith
      (closingFor site)
      (insideBrackets here (commaSep (fmap (align . hsExpr ctx) xs)))
  RecordCon{..} ->
    name ctx rcon_con
      <> breakOrNothing
      <> indent (braces (insideBrackets here (commaSep (fmap align (fields <> wildcard)))))
    where
      HsRecFields{..} = rcon_flds
      fields = fmap (at_ ctx (fieldBind ctx (at_ ctx (name ctx . foLabel)))) rec_flds
      wildcard = case rec_dotdot of
        Just l -> [at ctx l (const (txt ".."))]
        Nothing -> []
  RecordUpd{..} ->
    hsExpr ctx rupd_expr <> breakOrNothing <> indent (braces (insideBrackets here updates))
    where
      updates = case rupd_flds of
        RegularRecUpdFields{..} ->
          commaSep (fmap (align . at_ ctx (fieldBind ctx (at_ ctx (fieldOcc ctx)))) recUpdFields)
        OverloadedRecUpdFields{..} ->
          commaSep (fmap (align . at_ ctx (fieldBind ctx (at_ ctx labelChain))) olRecUpdFields)
      labelChain (FieldLabelStrings flss) = dotFields ctx (unLoc <$> flss)
  HsGetField{..} ->
    hsExpr ctx gf_expr <> txt "." <> at ctx gf_field (dotField ctx)
  HsProjection{..} -> parens (txt "." <> dotFields ctx proj_flds)
  ExprWithTySig _ x HsWC{hswc_body} ->
    align $
      hsExpr ctx x
        <> joinedBy "::"
        <> indent (hsSigType ctx hswc_body)
  ArithSeq _ _ range -> arithSeq ctx here (closingFor site) range
  HsTypedBracket (bracketAnn, _) e ->
    txt opener
      <> breakOrNothing
      <> insideBrackets here (hsExpr ctx e)
      <> breakOrNothing
      <> txt "||]"
    where
      opener = case bracketAnn of
        BracketNoE{} -> "[||"
        BracketHasE{} -> "[e||"
  HsUntypedBracket _ q -> quotation ctx here q
  HsTypedSplice _ (HsTypedSpliceExpr _ e) -> spliceTH ctx (Is #typed) e DollarSplice
  HsUntypedSplice _ splice -> untypedSplice ctx DollarSplice splice
  HsProc _ p e ->
    txt "proc"
      <> layoutFrom ctx (spanOf p) (breakOrSpace <> indent (hsPat ctx p) <> breakOrSpace)
      <> txt "->"
      <> attachBody (CmdTopBody ctx plainSite e)
  HsStatic _ e -> txt "static" <> breakOrSpace <> indent (hsExpr ctx e)
  HsPragE _ prag x -> case prag of
    HsPragSCC _ n ->
      txt "{-# SCC "
        <> outputable n
        <> txt " #-}"
        <> breakOrSpace
        <> nest (if siteInBlock site then 1 else 0) (hsExpr ctx x)
  HsEmbTy _ HsWC{hswc_body} -> txt "type" <> space <> hsType ctx hswc_body
  HsHole holeKind -> case holeKind of
    HoleVar n -> name ctx n
    HoleError -> error "Tilia: a nameless hole cannot come from a successful parse"
  HsForAll _ tele e -> forallTelescope ctx tele <> breakOrSpace <> hsExpr ctx e
  HsQual _ qs e ->
    at ctx qs (contextOf loneVariableExpr (hsExpr ctx) . fmap unbracketed)
      <> joinedBy "=>"
      <> hsExpr ctx e
  HsFunArr _ multAnn x y ->
    hsExpr ctx x
      <> space
      <> multiplicity (hsExpr ctx) multAnn
      <> joinedBy "->"
      <> case unLoc y of
        HsFunArr{} -> exprBody ctx plainSite (spanOf y) (unLoc y)
        _ -> hsExpr ctx y

-- | The body of an equation, as the enclosing construct will hand it on.
type BodyOf body b = Site -> LocatedA body -> b

-- | A body standing on its own, with the given bracing.
bodyIn :: BodyOf body b -> Bracing -> LocatedA body -> b
bodyIn mkBody bracing = mkBody (withBracing bracing plainSite)

-- | Is this expression a literal?
isLiteral :: LHsExpr GhcPs -> Bool
isLiteral e = case unLoc e of
  HsLit{} -> True
  HsOverLit{} -> True
  _ -> False

-- | A function applied to arguments.
application :: Ctx -> Site -> LHsExpr GhcPs -> LHsExpr GhcPs -> Doc
application ctx site f x =
  case placement of
    Normal ->
      mayBraceWhenFlat (siteBracing site) headAndInit
        <> indent (breakOrSpace <> hsExpr ctx lastArg)
    Hanging ->
      layoutFrom ctx initSpan (headAndInit MayBrace)
        <> attach Hanging (hsExpr ctx lastArg)
  where
    (func, args) = gatherArgs f (x :| [])
    initArgs = NE.init args
    lastArg = NE.last args
    initSpan = spanOf f <> (startOf <$> spanOf lastArg)
    placement
      | maybe False (writtenOnOneLine ctx) initSpan = exprHangs (unLoc lastArg)
      | otherwise = Normal
    headAndInit bracing =
      hsExprIn ctx site{siteApplicand = True, siteBracing = bracing} func
        <> foldMap
          ( nest (if placement == Hanging then 0 else 1)
              . (breakOrSpace <>)
              . hsExprIn ctx (withBracing bracing plainSite)
          )
          initArgs

-- | The head of an application, and its arguments in the order written.
gatherArgs ::
  LHsExpr GhcPs ->
  NonEmpty (LHsExpr GhcPs) ->
  (LHsExpr GhcPs, NonEmpty (LHsExpr GhcPs))
gatherArgs f known = case unLoc f of
  HsApp _ l r -> gatherArgs l (NE.cons r known)
  _ -> (f, known)

-- | A tuple, or a tuple section.
tuple :: Ctx -> Maybe Span -> ClosingIndent -> Boxity -> [HsTupArg GhcPs] -> Doc
tuple ctx here closing boxity args
  | any isMissing args = flat (brackets' (sepBy comma (fmap arg args)))
  | otherwise = brackets' (insideBrackets here (commaSep (fmap arg args)))
  where
    brackets' = boxityBrackets boxity closing
    isMissing = \case
      Missing _ -> True
      _ -> False
    arg =
      align . \case
        Present _ e -> hsExpr ctx e
        Missing _ -> mempty

-- | An arithmetic sequence, in whichever of its four forms.
arithSeq :: Ctx -> Maybe Span -> ClosingIndent -> ArithSeqInfo GhcPs -> Doc
arithSeq ctx here closing = \case
  From from -> wrap (hsExpr ctx from <> breakOrSpace <> txt "..")
  FromThen from next ->
    wrap (commaSep (fmap (hsExpr ctx) [from, next]) <> breakOrSpace <> txt "..")
  FromTo from to ->
    wrap (hsExpr ctx from <> breakOrSpace <> txt ".." <> space <> hsExpr ctx to)
  FromThenTo from next to ->
    wrap $
      commaSep (fmap (hsExpr ctx) [from, next])
        <> breakOrSpace
        <> txt ".."
        <> space
        <> hsExpr ctx to
  where
    wrap = bracketsWith closing . insideBrackets here

-- | One field of a record construction or update.
fieldBind ::
  (HasLoc l) =>
  Ctx ->
  (GenLocated l a -> Doc) ->
  HsFieldBind (GenLocated l a) (LHsExpr GhcPs) ->
  Doc
fieldBind ctx label HsFieldBind{..} =
  label hfbLHS
    <> includeUnless
      hfbPun
      ( space
          <> keywordAt ctx (spanOf =<< hfbAnn) "="
          <> attach placement (hsExpr ctx hfbRHS)
      )
  where
    placement
      | sameLine (spanOf hfbLHS) (spanOf hfbRHS) = exprHangs (unLoc hfbRHS)
      | otherwise = Normal

-- | One field name written after a dot.
dotField :: Ctx -> DotFieldOcc GhcPs -> Doc
dotField ctx = name ctx . fmap (mkVarUnqual . field_label) . dfoLabel

-- | A run of field names, joined by the dots between them.
dotFields :: Ctx -> NonEmpty (DotFieldOcc GhcPs) -> Doc
dotFields ctx = sepBy (txt ".") . fmap (dotField ctx) . NE.toList

-- | A chain of operators applied to expressions.
exprChain ::
  Ctx ->
  Site ->
  LHsExpr GhcPs ->
  LHsExpr GhcPs ->
  LHsExpr GhcPs ->
  Doc
exprChain ctx site x op y =
  renderExprChain ctx site (uncurry (associate (fixityOf ctx)) chain)
  where
    chain = flattenAround splitOpApp x op y

-- | Take an operator application apart, if that is what this is.
splitOpApp ::
  LHsExpr GhcPs ->
  Maybe (LHsExpr GhcPs, LHsExpr GhcPs, LHsExpr GhcPs)
splitOpApp e = case unLoc e of
  OpApp _ l o r -> Just (l, o, r)
  _ -> Nothing

-- | The fixity of an operator used in an expression.
fixityOf :: Ctx -> LHsExpr GhcPs -> Maybe Fixity
fixityOf ctx o =
  (spanOf o >>= (`Map.lookup` ctxCaptured ctx))
    <|> (operatorName o >>= operatorFixity ctx InTerms)

-- | Print a chain of operators once precedence has regrouped it.
renderExprChain :: Ctx -> Site -> OpChain (LHsExpr GhcPs) (LHsExpr GhcPs) -> Doc
renderExprChain ctx site = \case
  Operand e -> hsExprIn ctx site e
  chain@(Chain operands@(firstOne :| _) operators) ->
    layoutFrom ctx (chainSpan spanOf chain) $
      if trailing
        then laidOut MayBrace
        else
          mayBraceWhenFlat
            (if placement == Hanging then MayBrace else NoBrace)
            laidOut
    where
      placement =
        chainPlacement exprHangs (chainSpan spanOf firstOne) (NE.last operands)
      laidOut bracing =
        renderExprChain ctx (withBracing bracing site) firstOne
          <> foldr (piece bracing) mempty (links operands operators)
      piece bracing (previous, o, operand, isLast) rest' =
        let operandSite =
              withBracing (if isLast then siteBracing site else bracing) plainSite
            rendered = renderExprChain ctx operandSite operand
         in if trailing
              then
                space
                  <> hsExpr ctx o
                  <> attach
                    (tailPlacement isLast (chainSpan spanOf previous) o operand)
                    (rendered <> rest')
              else
                attachOperator
                  placement
                  (chainSpan spanOf previous)
                  (spanOf o)
                  (chainSpan spanOf operand)
                  (hsExpr ctx o)
                  rendered
                  <> rest'
      tailPlacement isLast before o operand
        | isLast,
          not (maybe True isSingleLine (chainSpan spanOf operand)),
          not (lineCommentWrittenBetween ctx before (spanOf o)) =
            chainPlacement exprHangs (spanOf o) operand
        | otherwise = Normal
      commentedOperators =
        or (zipWith commentedBefore (NE.toList operands) operators)
      commentedBefore operand o =
        commentPrintedBetween ctx (chainSpan spanOf operand) (spanOf o)
      trailing =
        (length operators == 1 || endsHanging)
          && not commentedOperators
          && and (zipWith couldTrail (NE.toList operands) operators)
      endsHanging = exprHangs (unLoc (lastOperand chain)) == Hanging
      couldTrail previous o =
        isSeparator (fixityOf ctx o)
          && maybe False (writtenOnOneLine ctx) (chainSpan spanOf previous)
          && placement == Normal
          && not (isDoBlock (lastOperand previous))

-- | Is this expression a @do@ or @mdo@ block?
isDoBlock :: LHsExpr GhcPs -> Bool
isDoBlock e = case unLoc e of
  HsDo _ (DoExpr _) _ -> True
  HsDo _ (MDoExpr _) _ -> True
  _ -> False

-- | Whether an operand hangs, which it can only where it starts on the line
-- the given span starts on.
chainPlacement ::
  (HasLoc l) =>
  (a -> Placement) ->
  Maybe Span ->
  OpChain (GenLocated l a) op ->
  Placement
chainPlacement placer leader operand = case operand of
  Operand (L _ n) | startsTogether -> placer n
  Chain (first :| _) _ -> chainPlacement placer leader first
  _ -> Normal
  where
    startsTogether =
      case (leader, chainSpan spanOf operand) of
        (Just a, Just b) -> spanStartLine a == spanStartLine b
        _ -> False

-- | Is this quoted constraint nothing but a variable?
loneVariableExpr :: LHsExpr GhcPs -> Bool
loneVariableExpr e = case unLoc e of
  HsVar _ (L _ n) -> isVarOcc (rdrNameOcc n)
  _ -> False

-- | What is inside, however many brackets were written around it.
unbracketed :: LHsExpr GhcPs -> LHsExpr GhcPs
unbracketed e = case unLoc e of
  HsPar _ inner -> unbracketed inner
  _ -> e

-- | A statement of a comprehension or a guard.
hsStmt :: Ctx -> ExprLStmt GhcPs -> Doc
hsStmt ctx = at_ ctx (stmtBody ctx plainSite (ExprBody ctx))

-- | One statement of a block, a comprehension or a guard.
stmtBody ::
  ( Body b,
    Anno [LStmt GhcPs (XRec GhcPs body)] ~ SrcSpanAnnLW,
    Anno (Stmt GhcPs (XRec GhcPs body)) ~ SrcSpanAnnA,
    Anno body ~ SrcSpanAnnA
  ) =>
  Ctx ->
  Site ->
  BodyOf body b ->
  Stmt GhcPs (XRec GhcPs body) ->
  Doc
stmtBody ctx site mkBody = \case
  LastStmt _ body _ _ -> printBody (mkBody site body)
  BodyStmt _ body _ _ -> printBody (mkBody site body)
  BindStmt arrow p f ->
    hsPatClosing ctx (closingFor site) p
      <> nest 1 (space <> keywordAt ctx (spanOf arrow) "<-")
      <> layoutFrom
        ctx
        (fmap endOf (spanOf p) <> spanOf f)
        (attach placement (printBody bound))
    where
      bound = mkBody plainSite f
      placement
        | sameLine (spanOf p) (spanOf f) = bodyPlacement bound
        | otherwise = Normal
  LetStmt _ binds -> txt "let" <> space <> align (bound binds)
    where
      bound = \case
        EmptyLocalBinds _ -> txt "{}"
        bs -> localBinds ctx (siteBracing site) bs
  ParStmt{} ->
    error "Tilia: ParStmt should have been unpacked"
  TransStmt{..} -> case (trS_form, trS_by) of
    (ThenForm, Nothing) ->
      txt "then" <> breakOrSpace <> indent (hsExpr ctx trS_using)
    (ThenForm, Just e) ->
      txt "then"
        <> breakOrSpace
        <> indent (hsExpr ctx trS_using)
        <> breakOrSpace
        <> txt "by"
        <> breakOrSpace
        <> indent (hsExpr ctx e)
    (GroupForm, Nothing) ->
      txt "then group using" <> breakOrSpace <> indent (hsExpr ctx trS_using)
    (GroupForm, Just e) ->
      txt "then group by"
        <> breakOrSpace
        <> indent (hsExpr ctx e)
        <> breakOrSpace
        <> txt "using"
        <> breakOrSpace
        <> indent (hsExpr ctx trS_using)
  RecStmt{..} ->
    txt "rec"
      <> space
      <> align
        ( at ctx recS_stmts $ \xs ->
            items (siteBracing site) $
              keepBlanks
                (separatedByBlank ctx)
                [ (spanOf s, at_ ctx (stmtBody ctx site mkBody) s)
                | s <- xs
                ]
        )

-- | The statements of a block, with the break that introduces them.
statements ::
  ( Body b,
    Anno [LStmt GhcPs (XRec GhcPs body)] ~ SrcSpanAnnLW,
    Anno (Stmt GhcPs (XRec GhcPs body)) ~ SrcSpanAnnA,
    Anno body ~ SrcSpanAnnA
  ) =>
  Ctx ->
  Site ->
  BodyOf body b ->
  XRec GhcPs [LStmt GhcPs (XRec GhcPs body)] ->
  Doc
statements ctx site mkBody es =
  breakOrSpace <> underSite site (at ctx es block)
  where
    block xs =
      items (siteBracing site) $
        keepBlanks
          (separatedByBlank ctx)
          [ (spanOf stmt, item position stmt)
          | (position, stmt) <- positioned xs
          ]
    item position stmt =
      blockItem position (siteBracing site) $ \bracing ->
        at_ ctx (stmtBody ctx (blockSite bracing) mkBody) stmt
    blockSite bracing =
      plainSite{siteInBlock = True, siteBracing = bracing}

-- | A list comprehension.
comprehension :: Ctx -> Site -> Maybe Span -> XRec GhcPs [ExprLStmt GhcPs] -> Doc
comprehension ctx site here es = align (variant onOneLine acrossLines)
  where
    onOneLine = txt "[" <> body <> txt "]"
    acrossLines = txt "[" <> space <> keepInside (body <> hardBreak <> txt "]")
    keepInside = if siteInBlock site then align else id
    body = insideBrackets here (at ctx es sections)
    sections xs = case unsnoc xs of
      Nothing -> error "Tilia: a comprehension always yields something"
      Just (stmts, yield) ->
        align (hsStmt ctx yield)
          <> breakOrSpace
          <> sepBy breakOrSpace (fmap section (comprehensionSections stmts))
    section stmts =
      located'
        (spansOf stmts)
        (txt "|" <> space <> align (commaSep (fmap (align . hsStmt ctx) stmts)))
    located' = maybe id located

-- | Split the statements of a comprehension into its parallel sections.
comprehensionSections :: [ExprLStmt GhcPs] -> [[ExprLStmt GhcPs]]
comprehensionSections = fmap unnest . branches
  where
    branches = \case
      [L _ (ParStmt _ blocks _ _)] ->
        [run | ParStmtBlock _ run _ _ <- NE.toList blocks]
      run -> [run]
    unnest = concatMap $ \case
      L _ ParStmt{} -> error "Tilia: parallel blocks do not nest"
      stmt@(L _ TransStmt{trS_stmts}) -> unnest trS_stmts <> [stmt]
      stmt -> [stmt]

-- | A @case@ expression or command.
caseOf ::
  ( Body b,
    Anno (GRHS GhcPs (LocatedA body)) ~ EpAnnCO,
    Anno (Match GhcPs (LocatedA body)) ~ SrcSpanAnnA
  ) =>
  Ctx ->
  Site ->
  BodyOf body b ->
  EpAnnHsCase ->
  LHsExpr GhcPs ->
  MatchGroup GhcPs (LocatedA body) ->
  Doc
caseOf ctx site mkBody EpAnnHsCase{..} scrutinee mg =
  keywordAt ctx (spanOf hsCaseAnnCase) "case"
    <> space
    <> hsExpr ctx scrutinee
    <> space
    <> indent (keywordAt ctx (spanOf hsCaseAnnOf) "of")
    <> breakOrSpace
    <> bracedForSite site alternatives
  where
    alternatives b = underSite site (matchGroup ctx b mkBody CaseStyle mg)

-- | A lambda, in any of its three spellings.
lambda ::
  ( Body b,
    Anno (GRHS GhcPs (LocatedA body)) ~ EpAnnCO,
    Anno (Match GhcPs (LocatedA body)) ~ SrcSpanAnnA
  ) =>
  Ctx ->
  Site ->
  HsLamVariant ->
  BodyOf body b ->
  MatchGroup GhcPs (LocatedA body) ->
  Doc
lambda ctx site variant' mkBody mg = case keyword of
  Nothing -> matchGroup ctx (siteBracing site) mkBody LambdaStyle mg
  Just kw -> txt kw <> breakOrSpace <> bracedForSite site alternatives
  where
    alternatives b = underSite site (matchGroup ctx b mkBody LambdaCaseStyle mg)
    keyword = case variant' of
      LamSingle -> Nothing
      LamCase -> Just "\\case"
      LamCases -> Just "\\cases"

-- | An @if@ expression or command.
ifThenElse ::
  (Body b) =>
  Ctx ->
  (LocatedA body -> b) ->
  AnnsIf ->
  LHsExpr GhcPs ->
  LocatedA body ->
  LocatedA body ->
  Doc
ifThenElse ctx bodyOf AnnsIf{aiThen, aiElse} condition thenBody elseBody =
  txt "if"
    <> space
    <> hsExpr ctx condition
    <> breakOrSpace
    <> indent
      ( branch (locA aiThen) "then" thenBody
          <> breakOrSpace
          <> branch (locA aiElse) "else" elseBody
      )
  where
    branch written word body =
      keywordAt ctx keywordSpan word
        <> space
        <> layoutFrom
          ctx
          (keywordSpan <> spanOf body)
          (attach placement (printBody (bodyOf body)))
      where
        keywordSpan = spanOfSrcSpan written
        placement
          | commentPrintedBetween ctx keywordSpan (spanOf body) = Normal
          | otherwise = bodyPlacement (bodyOf body)

-- | A @let@ expression or command, with @in@ ending its line where it is
-- one of a chain written that way, and the next on the line under it.
letIn ::
  (Body b) =>
  Ctx ->
  (LocatedA body -> b) ->
  -- | Whether the body is another @let@.
  Bool ->
  (EpToken "let", EpToken "in") ->
  HsLocalBinds GhcPs ->
  LocatedA body ->
  Doc
letIn ctx bodyOf nested (letToken, inToken) binds body =
  align $
    written (spanOf letToken) (txt "let")
      <> space
      <> align
        ( localBinds ctx NoBrace binds
            <> foldMap (emptyAnchor . startOf) (spanOf inToken)
        )
      <> written (spanOf inToken) (separator <> txt "in")
      <> (if chained then breakOrSpace else space)
      <> align (printBody (bodyOf body))
  where
    chained = nested && not (sameLine (spanOf inToken) (spanOf body))
    separator = if chained then space else variant space (hardBreak <> txt " ")
    written s d = if chained then fenceWithin ctx s d else d

-- | The bindings of a @let@ or a @where@.
localBinds :: Ctx -> Bracing -> HsLocalBinds GhcPs -> Doc
localBinds ctx bracing = \case
  HsValBinds ann (ValBinds _ binds sigs) ->
    anchored ann . align . items bracing $
      keepBlanks
        (separatedByBlank ctx)
        [ (spanOf item, rendered position item)
        | (position, item) <- positioned sorted
        ]
    where
      sorted =
        sortBy
          (leftmost_smallest `on` getLocA)
          (fmap (fmap Left) binds <> fmap (fmap Right) sigs)
      rendered position item =
        blockItem position NoBrace $ \b ->
          at_ ctx (either (valDecl ctx b) (knotSig (ctxKnot ctx) ctx)) item
  HsValBinds _ _ -> error "Tilia: renamer-only local bindings"
  HsIPBinds ann (IPBinds _ xs) ->
    anchored ann (items bracing (fmap (at_ ctx implicitBind) xs))
  EmptyLocalBinds _ -> mempty
  where
    implicitBind (IPBind _ (L _ n) e) =
      outputable n
        <> joinedBy "="
        <> indent (hsExprIn ctx (withBracing MayBrace plainSite) e)
    anchored ann d = case ann of
      EpAnn{anns = AnnList{al_anchor}}
        | not (isZeroWidthSpan (locA al_anchor)) ->
            atSpan ctx (spanOfSrcSpan (locA al_anchor)) d
      _ -> d

-- | Where a @do@ or @mdo@ was written.
doKeywordSpan :: AnnList EpaLocation -> Maybe Span
doKeywordSpan = spanOfSrcSpan . locA . al_rest

-- | Where the @where@ keyword of a group of local bindings was.
whereKeywordSpan :: HsLocalBinds GhcPs -> Maybe Span
whereKeywordSpan =
  spanOfSrcSpan . \case
    HsValBinds EpAnn{anns = AnnList{al_rest}} _ -> locA al_rest
    HsIPBinds EpAnn{anns = AnnList{al_rest}} _ -> locA al_rest
    EmptyLocalBinds _ -> noSrcSpan

-- | Does this @let@ or @where@ bind nothing at all?
isEmptyLocalBinds :: HsLocalBinds GhcPs -> Bool
isEmptyLocalBinds = \case
  EmptyLocalBinds _ -> True
  HsValBinds _ (ValBinds _ binds sigs) -> null binds && null sigs
  _ -> False

-- | An arrow-notation command.
hsCmd :: Ctx -> Site -> LHsCmd GhcPs -> Doc
hsCmd ctx site l = at ctx l (cmdBody ctx site)

-- | Every shape an arrow-notation command can take.
cmdBody :: Ctx -> Site -> HsCmd GhcPs -> Doc
cmdBody ctx site = \case
  HsCmdArrApp _ body input arrow rightToLeft ->
    let writtenFirst = if rightToLeft then body else input
        writtenSecond = if rightToLeft then input else body
     in hsExprIn ctx site{siteApplicand = False} writtenFirst
          <> breakOrSpace
          <> indent
            ( txt (arrowText arrow rightToLeft)
                <> attach (exprHangs (unLoc input)) (hsExpr ctx writtenSecond)
            )
  HsCmdArrForm _ form Prefix cmds ->
    bananaWith (closingFor site) $
      hsExpr ctx form
        <> includeUnless
          (null cmds)
          ( breakOrSpace
              <> indent
                ( sepBy
                    breakOrSpace
                    (fmap (printBody . CmdTopBody ctx plainSite) cmds)
                )
          )
  HsCmdArrForm _ form Infix [l, r] -> cmdChain ctx site l form r
  HsCmdArrForm _ _ Infix _ ->
    error "Tilia: an infix command form always has exactly two operands"
  HsCmdApp _ cmd e ->
    hsCmd ctx site{siteApplicand = True} cmd
      <> breakOrSpace
      <> indent (hsExpr ctx e)
  HsCmdLam _ variant' mg -> lambda ctx site variant' (CmdBody ctx) mg
  HsCmdPar _ c -> parens (hsCmd ctx plainSite c)
  HsCmdCase anns e mg -> caseOf ctx site (CmdBody ctx) anns e mg
  HsCmdIf anns _ c t e ->
    ifThenElse ctx (bodyIn (CmdBody ctx) (siteBracing site)) anns c t e
  HsCmdLet tokens binds c ->
    letIn ctx (bodyIn (CmdBody ctx) (siteBracing site)) nested tokens binds c
    where
      nested = case unLoc c of
        HsCmdLet{} -> True
        _ -> False
  HsCmdDo anns es ->
    keywordAt ctx (doKeywordSpan anns) "do"
      <> statements ctx site (CmdBody ctx) es

-- | How the arrow of a command application is spelled.
arrowText :: HsArrAppType -> Bool -> Text
arrowText arrow rightToLeft = case (arrow, rightToLeft) of
  (HsFirstOrderApp, True) -> "-<"
  (HsHigherOrderApp, True) -> "-<<"
  (HsFirstOrderApp, False) -> ">-"
  (HsHigherOrderApp, False) -> ">>-"

-- | A command at the top of an arrow form.
cmdTop :: Ctx -> Site -> HsCmdTop GhcPs -> Doc
cmdTop ctx site (HsCmdTop _ cmd) = hsCmd ctx site cmd

-- | A chain of operators applied to commands.
cmdChain ::
  Ctx ->
  Site ->
  LHsCmdTop GhcPs ->
  LHsExpr GhcPs ->
  LHsCmdTop GhcPs ->
  Doc
cmdChain ctx site l op r =
  render (uncurry (associate (fixityOf ctx)) (flattenAround splitCmd l op r))
  where
    splitCmd c = case unLoc c of
      HsCmdTop _ (L _ (HsCmdArrForm _ o Infix [a, b])) -> Just (a, o, b)
      _ -> Nothing
    render = \case
      Operand c -> at ctx c (cmdTop ctx site)
      chain@(Chain operands@(firstOne :| _) operators) ->
        layoutFrom ctx (chainSpan spanOf chain) $
          mayBraceWhenFlat (if placement == Hanging then MayBrace else NoBrace) laidOut
        where
          placement =
            chainPlacement
              cmdTopHangs
              (chainSpan spanOf firstOne)
              (NE.last operands)
          laidOut bracing =
            renderIn bracing firstOne
              <> foldMap (piece bracing) (links operands operators)
          piece bracing (previous, o, operand, isLast) =
            attachOperator
              placement
              (chainSpan spanOf previous)
              (spanOf o)
              (chainSpan spanOf operand)
              (hsExpr ctx o)
              (renderIn (if isLast then siteBracing site else bracing) operand)
          renderIn bracing = \case
            Operand c -> at ctx c (cmdTop ctx (withBracing bracing site))
            inner -> render inner

-- | A value binding.
valDecl :: Ctx -> Bracing -> HsBind GhcPs -> Doc
valDecl ctx bracing = \case
  FunBind _ funId funMatches ->
    matchGroup ctx bracing (ExprBody ctx) (FunctionStyle funId) funMatches
  PatBind _ p multAnn grhss ->
    match ctx bracing (ExprBody ctx) PatternBindStyle False multAnn NoSrcStrict [p] grhss
  PatSynBind _ psb -> patSynBind ctx psb
  VarBind{} -> error "Tilia: VarBind is introduced by the type checker"

-- | Which shape a group of equations takes.
data MatchStyle
  = -- | @f x = …@
    FunctionStyle (LocatedN RdrName)
  | -- | @(x, y) = …@
    PatternBindStyle
  | -- | An alternative of a @case@
    CaseStyle
  | -- | The body of a @\\@
    LambdaStyle
  | -- | An alternative of a @\\case@ or @\\cases@
    LambdaCaseStyle

-- | What separates a guard from what it guards.
data GuardStyle
  = -- | @| p = …@
    EqualsSign
  | -- | @| p -> …@
    RightArrow
  deriving (Eq, Show)

-- | A group of equations.
matchGroup ::
  ( Body b,
    Anno (GRHS GhcPs (LocatedA body)) ~ EpAnnCO,
    Anno (Match GhcPs (LocatedA body)) ~ SrcSpanAnnA
  ) =>
  Ctx ->
  Bracing ->
  BodyOf body b ->
  MatchStyle ->
  MatchGroup GhcPs (LocatedA body) ->
  Doc
matchGroup ctx bracing mkBody style MG{..}
  | isCaseStyle style, null (unLoc mg_alts) = txt "{}"
  | otherwise = items blockBracing (fmap rendered (positioned (unLoc mg_alts)))
  where
    blockBracing = case style of
      CaseStyle -> ifEmpty
      LambdaCaseStyle -> ifEmpty
      _ -> NoBrace
    ifEmpty = if null (unLoc mg_alts) then MayBrace else bracing
    rendered (position, m) =
      blockItem position bracing (\b -> at_ ctx (renderMatch b) m)
    renderMatch b m@Match{..} =
      match
        ctx
        b
        mkBody
        (adjustStyle m style)
        (isInfixMatch m)
        (HsUnannotated EpPatBind)
        (bangBeforeName m)
        (unLoc m_pats)
        m_grhss

-- | Put this equation's own name into the style it is printed with. A
-- style that carries no name is left as it is.
adjustStyle :: Match GhcPs body -> MatchStyle -> MatchStyle
adjustStyle m = \case
  FunctionStyle _ | FunRhs{mc_fun = f} <- m_ctxt m -> FunctionStyle f
  style -> style

-- | Was a @!@ written in front of the name this equation defines?
bangBeforeName :: Match id body -> SrcStrictness
bangBeforeName = \case
  Match{m_ctxt = FunRhs{mc_strictness}} -> mc_strictness
  _ -> NoSrcStrict

-- | One equation: a head, a body, and possibly a @where@.
match ::
  (Body b, Anno (GRHS GhcPs (LocatedA body)) ~ EpAnnCO) =>
  Ctx ->
  Bracing ->
  BodyOf body b ->
  MatchStyle ->
  -- | Written infix?
  Bool ->
  HsMultAnn GhcPs ->
  SrcStrictness ->
  [LPat GhcPs] ->
  GRHSs GhcPs (LocatedA body) ->
  Doc
match ctx bracing mkBody style isInfix multAnn strict pats GRHSs{..} =
  multiplicity (hsType ctx) multAnn
    <> multAnnGap
    <> strictness strict
    <> head'
    <> nest bodyNest rightHandSide
    <> whereClause
  where
    bodyNest = if indentBody then 1 else 0
    rightHandSide = separator <> layoutFrom ctx bodySpan (attach placement body)

    multAnnGap = case multAnn of
      HsUnannotated{} -> mempty
      _ -> space

    indentBody = case pats of
      [] -> False
      _ ->
        not (maybe True isSingleLine headSpan)
          && not (isCaseStyle style && any containsOrPat pats)

    headSpan = case style of
      FunctionStyle n -> spanOf n <> patSpans
      _ -> patSpans
    patSpans = spansOf pats

    head' = case pats of
      [] -> case style of
        FunctionStyle n -> name ctx n
        _ -> mempty
      (headPat : tailPats) -> layoutFrom ctx headSpan $ case style of
        FunctionStyle n -> defHead (fromBool isInfix) (fromBool indentBody) (name ctx n) rendered
        PatternBindStyle -> sepBy breakOrSpace rendered
        CaseStyle -> sepBy breakOrSpace rendered
        LambdaStyle -> txt "\\" <> lambdaGap headPat <> align (sepBy breakOrSpace rendered)
        LambdaCaseStyle ->
          hsPat ctx headPat
            <> includeUnless
              (null tailPats)
              (breakOrSpace <> indent (sepBy breakOrSpace (fmap (hsPat ctx) tailPats)))
      where
        rendered = fmap (hsPat ctx) pats

    lambdaGap p = includeWhen (needsGap (unLoc p)) space
    needsGap = \case
      LazyPat{} -> True
      BangPat{} -> True
      SplicePat{} -> True
      InvisPat{} -> True
      _ -> False

    endOfPats = case pats of
      [] -> case style of
        FunctionStyle n -> spanOf n
        _ -> Nothing
      _ -> spanOf (last pats)

    hasGuards = any (not . null . guardsOf . unLoc) grhssGRHSs

    rhsSpan = foldr1 (<>) (fmap (grhsSpan . unLoc) grhssGRHSs)
    bodySpan = fmap endOf endOfPats <> rhsSpan

    placement = case endOfPats of
      Just spn
        | any (longGuard . unLoc) grhssGRHSs || not (sameLine (Just spn) rhsSpan) ->
            Normal
      _ -> blockPlacement (bodyIn mkBody bracing) grhssGRHSs
    longGuard grhs = case guardsOf grhs of
      [] -> False
      [g] -> not (maybe True isSingleLine (spanOf g))
      _ -> True

    separator
      | length grhssGRHSs > 1 = mempty
      | otherwise = case style of
          FunctionStyle _ | hasGuards -> mempty
          FunctionStyle _ -> space <> indent (txt "=")
          PatternBindStyle | hasGuards -> mempty
          PatternBindStyle -> space <> indent (txt "=")
          s | isCaseStyle s && hasGuards -> mempty
          _ -> space <> txt "->"

    body = sepBy breakOrSpace (fmap alternative (NE.toList grhssGRHSs))

    alternative g =
      fenceWithin ctx (spanOf g) $
        guardedRhs ctx placement bracing mkBody groupStyle g

    groupStyle
      | isCaseStyle style && hasGuards = RightArrow
      | otherwise = EqualsSign

    whereClause = case grhssLocalBinds of
      EmptyLocalBinds _ -> mempty
      binds ->
        nest bodyNest . indent $
          breakOrSpace
            <> keywordAt ctx (whereKeywordSpan binds) "where"
            <> includeUnless
              (isEmptyLocalBinds binds)
              (breakOrSpace <> indent (localBinds ctx bracing binds))

-- | Is this an alternative of a @case@ or a @\\case@?
isCaseStyle :: MatchStyle -> Bool
isCaseStyle = \case
  CaseStyle -> True
  LambdaCaseStyle -> True
  _ -> False

-- | Does this pattern hold an or-pattern anywhere inside it?
containsOrPat :: LPat GhcPs -> Bool
containsOrPat = any isOrPat . listify (const True :: Pat GhcPs -> Bool)
  where
    isOrPat = \case
      OrPat{} -> True
      _ -> False

-- | Where an alternative was written, guards and all.
grhsSpan :: GRHS GhcPs (LocatedA body) -> Maybe Span
grhsSpan (GRHS _ guards body) = spanOf body <> spansOf guards

-- | The guards of an alternative.
guardsOf :: GRHS GhcPs body -> [GuardLStmt GhcPs]
guardsOf (GRHS _ guards _) = guards

-- | The placement of a body that is the whole of an equation.
blockPlacement ::
  (Body b) =>
  (LocatedA body -> b) ->
  NonEmpty (LGRHS GhcPs (LocatedA body)) ->
  Placement
blockPlacement bodyOf = \case
  L _ (GRHS _ _ body) :| [] -> bodyPlacement (bodyOf body)
  _ -> Normal

-- | One alternative of an equation: its guards, and what they guard.
guardedRhs ::
  (Body b, Anno (GRHS GhcPs (LocatedA body)) ~ EpAnnCO) =>
  Ctx ->
  -- | How the equation as a whole is placed.
  Placement ->
  -- | Bracing the body inherits.
  Bracing ->
  BodyOf body b ->
  GuardStyle ->
  LGRHS GhcPs (LocatedA body) ->
  Doc
guardedRhs ctx parentPlacement bracing mkBody style lgrhs@(L _ grhs) = case guards of
  [] -> atSpan ctx (grhsSpan grhs) (printBody bound)
  _ ->
    maybe id located (spanOf lgrhs) . layoutFrom ctx (grhsSpan grhs) $
      fenceWithin ctx headSpan guardHead
        <> nest
          (if parentPlacement == Normal then 1 else 0)
          (attach placement (printBody bound))
  where
    GRHS EpAnn{anns = GrhsAnn{ga_vbar, ga_sep}} guards body = grhs
    guardHead =
      txt "|"
        <> space
        <> align (commaSep (fmap (align . hsStmt ctx) guards))
        <> space
        <> indent (txt separator)
    headSpan = (<>) <$> (spanOf =<< ga_vbar) <*> separatorSpan
    bound = bodyIn mkBody bracing body
    separator = case style of
      EqualsSign -> "="
      RightArrow -> "->"
    placement
      | sameLine separatorSpan (spanOf body),
        not (commentPrintedBetween ctx (spansOf guards) separatorSpan),
        not (remarkBefore && remarkAfter) =
          bodyPlacement bound
      | otherwise = Normal
    separatorSpan = either spanOf spanOf ga_sep
    remarkBefore = lineCommentWrittenBetween ctx (spansOf guards) separatorSpan
    remarkAfter =
      lineCommentWrittenBetween ctx separatorSpan (nextLine <$> separatorSpan)
    nextLine s = mkSpan (spanEndLine s + 1, 1) (spanEndLine s + 1, 1)

-- | A pattern synonym binding.
patSynBind :: Ctx -> PatSynBind GhcPs GhcPs -> Doc
patSynBind ctx PSB{..} =
  txt "pattern" <> case psb_args of
    PrefixCon args ->
      space
        <> name ctx psb_id
        <> indent
          ( layoutAcross ctx args (argsAfterName (fmap (name ctx) args) (null args))
              <> definition (spansOf args)
          )
    RecCon args ->
      space
        <> name ctx psb_id
        <> indent
          ( layoutAcross
              ctx
              (vars args)
              ( includeUnless (null args) breakOrNothing
                  <> braces (commaSep (fmap (name ctx) (vars args)))
              )
              <> definition (spansOf (vars args))
          )
      where
        vars = fmap recordPatSynPatVar
    InfixCon l r ->
      layoutFrom
        ctx
        (spanOf l <> spanOf r)
        ( space
            <> name ctx l
            <> attachOperator
              Normal
              (spanOf l)
              (spanOf psb_id)
              (spanOf r)
              (name ctx psb_id)
              (name ctx r)
        )
        <> indent (definition (spanOf l <> spanOf r))
  where
    argsAfterName rendered isEmpty =
      includeUnless isEmpty breakOrSpace <> align (sepBy breakOrSpace rendered)

    definition argSpans =
      space <> case psb_dir of
        Unidirectional -> rhs "<-"
        ImplicitBidirectional -> rhs "="
        ExplicitBidirectional mg ->
          rhs "<-"
            <> breakOrSpace
            <> txt "where"
            <> breakOrSpace
            <> indent (matchGroup ctx NoBrace (ExprBody ctx) (FunctionStyle psb_id) mg)
      where
        rhs arrow =
          layoutFrom ctx (spanOf psb_id <> spanOf psb_def <> argSpans) $
            txt arrow <> breakOrSpace <> hsPat ctx psb_def

-- | An untyped splice, either @$x@ or a quasi-quotation.
untypedSplice :: Ctx -> SpliceDecoration -> HsUntypedSplice GhcPs -> Doc
untypedSplice ctx deco = \case
  HsUntypedSpliceExpr _ e -> spliceTH ctx (Isn't #typed) e deco
  HsQuasiQuote _ quoter str ->
    txt "["
      <> name ctx quoter
      <> txt "|"
      <> verbatim (T.pack (unpackFS (unLoc str)))
      <> txt "|]"

-- | A Template Haskell splice, typed or not.
spliceTH :: Ctx -> Choice "typed" -> LHsExpr GhcPs -> SpliceDecoration -> Doc
spliceTH ctx typed e = \case
  DollarSplice -> txt (if isTrue typed then "$$" else "$") <> spliced
  BareSplice -> spliced
  where
    spliced = at ctx e (align . exprBody ctx plainSite (spanOf e))

-- | A Template Haskell quotation.
quotation :: Ctx -> Maybe Span -> HsQuote GhcPs -> Doc
quotation ctx here = \case
  ExpBr (bracketAnn, _) e -> quoted (flavour bracketAnn) (hsExpr ctx e)
    where
      flavour = \case
        BracketNoE{} -> ""
        BracketHasE{} -> "e"
  PatBr _ p -> quoted "p" (hsPat ctx p)
  DecBrL _ decls ->
    quoted "d" (starGuard decls (knotDecls (ctxKnot ctx) ctx Free decls))
  DecBrG _ _ -> error "Tilia: DecBrG is produced by the renamer"
  TypBr _ ty -> quoted "t" (starGuard ty (hsType ctx ty))
  VarBr _ isSingle n -> txt (if isSingle then "'" else "''") <> name ctx n
  where
    quoted flavour body =
      txt "["
        <> txt flavour
        <> txt "|"
        <> breakOrNothing
        <> indent (insideBrackets here body <> breakOrNothing <> txt "|]")
    starGuard x body
      | risky = space <> body <> space
      | otherwise = body
      where
        risky =
          any isStar (listify (const True :: HsType GhcPs -> Bool) x)
            || any isAbstract (listify (const True :: FamilyInfo GhcPs -> Bool) x)
    isStar = \case
      HsStarTy{} -> True
      _ -> False
    isAbstract = \case
      ClosedTypeFamily Nothing -> True
      _ -> False

-- | An item of a block: the last one keeps the block's bracing, and any
-- other may brace itself when the block is flat.
blockItem :: Position -> Bracing -> (Bracing -> Doc) -> Doc
blockItem position bracing render = case position of
  Last -> render bracing
  Only -> render bracing
  _ -> mayBraceWhenFlat bracing render

-- | Let a block brace itself when the surrounding layout is flat.
mayBraceWhenFlat :: Bracing -> (Bracing -> Doc) -> Doc
mayBraceWhenFlat whenBroken render =
  variant (render MayBrace) (render whenBroken)

-- | A @case@ or lambda standing as a block item has to delimit itself when
-- the block is flat, unless it is the function of an application, where the
-- argument that follows already does so.
bracedForSite :: Site -> (Bracing -> Doc) -> Doc
bracedForSite site render
  | siteInBlock site && not (siteApplicand site) =
      mayBraceWhenFlat (siteBracing site) render
  | otherwise = render (siteBracing site)

----------------------------------------------------------------------------
-- Bodies

-- | An expression standing as the body of an enclosing construct.
data ExprBody = ExprBody Ctx Site (LHsExpr GhcPs)

instance Body ExprBody where
  printBody (ExprBody ctx site e) = hsExprIn ctx site e
  bodyPlacement (ExprBody _ _ e) = exprHangs (unLoc e)

-- | A command standing as the body of an enclosing construct.
data CmdBody = CmdBody Ctx Site (LHsCmd GhcPs)

instance Body CmdBody where
  printBody (CmdBody ctx site c) = hsCmd ctx site c
  bodyPlacement (CmdBody _ _ c) = cmdHangs (unLoc c)

-- | A command at the top of an arrow form.
data CmdTopBody = CmdTopBody Ctx Site (LHsCmdTop GhcPs)

instance Body CmdTopBody where
  printBody (CmdTopBody ctx site l) = at ctx l (cmdTop ctx site)
  bodyPlacement (CmdTopBody _ _ l) = cmdTopHangs (unLoc l)

-- | Does this expression absorb the line break that introduces it?
exprHangs :: HsExpr GhcPs -> Placement
exprHangs = \case
  HsDo _ (DoExpr _) _ -> Hanging
  HsDo _ (MDoExpr _) _ -> Hanging
  HsCase{} -> Hanging
  HsLam _ lamVariant mg -> case lamVariant of
    LamCase -> Hanging
    LamCases -> Hanging
    LamSingle -> case mg of
      MG _ (L _ [L _ (Match _ _ (L _ ps@(_ : _)) _)])
        | maybe False isSingleLine (spansOf ps) -> Hanging
      _ -> Normal
  HsProc _ p _
    | maybe False isSingleLine (spanOf p) -> Hanging
    | otherwise -> Normal
  HsApp _ _ y -> exprHangs (unLoc y)
  OpApp _ x op y
    | Just n <- operatorName op,
      occNameString (rdrNameOcc n) == "$" ->
        exprHangs (unLoc y)
    | otherwise -> exprHangs (unLoc x)
  _ -> Normal

-- | Does this command absorb the line break that introduces it?
cmdHangs :: HsCmd GhcPs -> Placement
cmdHangs = \case
  HsCmdDo{} -> Hanging
  HsCmdCase{} -> Hanging
  HsCmdLam{} -> Hanging
  _ -> Normal

cmdTopHangs :: HsCmdTop GhcPs -> Placement
cmdTopHangs (HsCmdTop _ c) = cmdHangs (unLoc c)

-- | The name of an operator, when the expression standing as one is a name.
operatorName :: LHsExpr GhcPs -> Maybe RdrName
operatorName e = case unLoc e of
  HsVar _ (L _ n) -> Just n
  _ -> Nothing