Data.hs

{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecordWildCards #-}

-- | Renedring of data type declarations.
module Ormolu.Printer.Meat.Declaration.Data
  ( p_dataDecl,
  )
where

import Control.Monad
import Data.List.NonEmpty (NonEmpty (..))
import Data.List.NonEmpty qualified as NE
import Data.Maybe (isJust, maybeToList)
import Data.Void
import GHC.Data.Strict qualified as Strict
import GHC.Hs
import GHC.Types.Fixity
import GHC.Types.ForeignCall
import GHC.Types.Name.Reader
import GHC.Types.SrcLoc
import Ormolu.Printer.Combinators
import Ormolu.Printer.Meat.Common
import Ormolu.Printer.Meat.Type
import Ormolu.Utils

p_dataDecl ::
  -- | Whether to format as data family
  FamilyStyle ->
  -- | Type constructor
  LocatedN RdrName ->
  -- | Type variables
  [tyVar] ->
  -- | Get location information for type variables
  (tyVar -> SrcSpan) ->
  -- | How to print type variables
  (tyVar -> R ()) ->
  -- | Lexical fixity
  LexicalFixity ->
  -- | Data definition
  HsDataDefn GhcPs ->
  R ()
p_dataDecl style name tyVars getTyVarLoc p_tyVar fixity HsDataDefn {..} = do
  txt $ case dd_cons of
    NewTypeCon _ -> "newtype"
    DataTypeCons False _ -> "data"
    DataTypeCons True _ -> "type data"
  txt $ case style of
    Associated -> mempty
    Free -> " instance"
  let constructorSpans = getLocA name : fmap getTyVarLoc tyVars
      sigSpans = maybeToList . fmap getLocA $ dd_kindSig
      contextSpans = maybeToList . fmap getLocA $ dd_ctxt
      ctypeSpans = maybeToList . fmap getLocA $ dd_cType
      declHeaderSpans =
        constructorSpans ++ sigSpans ++ contextSpans ++ ctypeSpans
  switchLayout declHeaderSpans . inci $ do
    case unLoc <$> dd_cType of
      Nothing -> pure ()
      Just (CType prag header (type_, _)) -> do
        breakpoint
        p_sourceText prag
        case header of
          Nothing -> pure ()
          Just (Header h _) -> space *> p_sourceText h
        space
        p_sourceText type_
        txt " #-}"
    breakpoint
    forM_ dd_ctxt p_lhsContext
    switchLayout constructorSpans $
      p_infixDefHelper
        (isInfix fixity)
        True
        (p_rdrName name)
        (p_tyVar <$> tyVars)
    forM_ dd_kindSig $ \k -> do
      space
      txt "::"
      breakpoint
      inci $ located k p_hsType
  let dd_cons' = case dd_cons of
        NewTypeCon a -> [a]
        DataTypeCons _ as -> as
      gadt = isJust dd_kindSig || any (isGadt . unLoc) dd_cons'
  case dd_cons' of
    [] -> pure ()
    first_dd_cons : _ ->
      if gadt
        then inci $ do
          switchLayout declHeaderSpans $ do
            breakpoint
            txt "where"
          breakpoint
          sepSemi (located' (p_conDecl False)) dd_cons'
        else switchLayout (getLocA name : (getLocA <$> dd_cons')) . inci $ do
          let singleConstRec = isSingleConstRec dd_cons'
              compactLayoutAroundEquals =
                onTheSameLine
                  (getLocA name)
                  (combineSrcSpans' (conDeclConsSpans (unLoc first_dd_cons)))
              conDeclConsSpans = \case
                ConDeclGADT {..} -> getLocA <$> con_names
                ConDeclH98 {..} -> getLocA con_name :| []
          if hasHaddocks dd_cons'
            then newline
            else
              if singleConstRec && compactLayoutAroundEquals
                then space
                else breakpoint
          equals
          space
          layout <- getLayout
          let s =
                if layout == MultiLine || hasHaddocks dd_cons'
                  then newline >> txt "|" >> space
                  else space >> txt "|" >> space
              sitcc' =
                if hasHaddocks dd_cons' || not singleConstRec
                  then sitcc
                  else id
          sep s (sitcc' . located' (p_conDecl singleConstRec)) dd_cons'
  unless (null dd_derivs) breakpoint
  inci $ sep newline (located' p_hsDerivingClause) dd_derivs

p_conDecl ::
  Bool ->
  ConDecl GhcPs ->
  R ()
p_conDecl singleConstRec = \case
  ConDeclGADT {..} -> do
    mapM_ (p_hsDoc Pipe True) con_doc
    let conDeclSpn =
          fmap getLocA (NE.toList con_names)
            <> [getLocA con_bndrs]
            <> maybeToList (fmap getLocA con_mb_cxt)
            <> conArgsSpans
          where
            conArgsSpans = case con_g_args of
              PrefixConGADT xs -> getLocA . hsScaledThing <$> xs
              RecConGADT x _ -> [getLocA x]
    switchLayout conDeclSpn $ do
      let c :| cs = con_names
      p_rdrName c
      unless (null cs) . inci $ do
        commaDel
        sep commaDel p_rdrName cs
      inci $ do
        let conTy = case con_g_args of
              PrefixConGADT xs ->
                let go (HsScaled a b) t = addCLocAA t b (HsFunTy EpAnnNotUsed a b t)
                 in foldr go con_res_ty xs
              RecConGADT r _ ->
                addCLocAA r con_res_ty $
                  HsFunTy
                    EpAnnNotUsed
                    (HsUnrestrictedArrow noHsUniTok)
                    (la2la $ HsRecTy EpAnnNotUsed <$> r)
                    con_res_ty
            qualTy = case con_mb_cxt of
              Nothing -> conTy
              Just qs ->
                addCLocAA qs conTy $
                  HsQualTy NoExtField qs conTy
            quantifiedTy =
              addCLocAA con_bndrs qualTy $
                hsOuterTyVarBndrsToHsType (unLoc con_bndrs) qualTy
        space
        txt "::"
        if hasDocStrings (unLoc con_res_ty)
          then newline
          else breakpoint
        located quantifiedTy p_hsType
  ConDeclH98 {..} -> do
    mapM_ (p_hsDoc Pipe True) con_doc
    let conNameSpn = getLocA con_name
        conNameWithContextSpn =
          [ RealSrcSpan real Strict.Nothing
            | Just (EpaSpan real _) <- matchAddEpAnn AnnForall <$> epAnnAnns con_ext
          ]
            <> fmap getLocA con_ex_tvs
            <> maybeToList (fmap getLocA con_mb_cxt)
            <> [conNameSpn]
        conDeclSpn = conNameSpn : conArgsSpans
          where
            conArgsSpans = case con_args of
              PrefixCon [] xs -> getLocA . hsScaledThing <$> xs
              PrefixCon (v : _) _ -> absurd v
              RecCon l -> [getLocA l]
              InfixCon x y -> getLocA . hsScaledThing <$> [x, y]
    switchLayout conNameWithContextSpn $ do
      when con_forall $ do
        p_forallBndrs ForAllInvis p_hsTyVarBndr con_ex_tvs
        breakpoint
      forM_ con_mb_cxt p_lhsContext
    switchLayout conDeclSpn $ case con_args of
      PrefixCon [] xs -> do
        p_rdrName con_name
        let args = hsScaledThing <$> xs
            argsHaveDocs = conArgsHaveHaddocks args
            delimiter = if argsHaveDocs then newline else breakpoint
        unless (null xs) delimiter
        inci . sitcc $
          sep delimiter (sitcc . located' p_hsType) args
      PrefixCon (v : _) _ -> absurd v
      RecCon l -> do
        p_rdrName con_name
        breakpoint
        inciIf (not singleConstRec) (located l p_conDeclFields)
      InfixCon (HsScaled _ x) (HsScaled _ y) -> do
        located x p_hsType
        breakpoint
        inci $ do
          p_rdrName con_name
          space
          located y p_hsType

p_lhsContext ::
  LHsContext GhcPs ->
  R ()
p_lhsContext = \case
  L _ [] -> pure ()
  ctx -> do
    located ctx p_hsContext
    space
    txt "=>"
    breakpoint

isGadt :: ConDecl GhcPs -> Bool
isGadt = \case
  ConDeclGADT {} -> True
  ConDeclH98 {} -> False

p_hsDerivingClause ::
  HsDerivingClause GhcPs ->
  R ()
p_hsDerivingClause HsDerivingClause {..} = do
  txt "deriving"
  let derivingWhat = located deriv_clause_tys $ \case
        DctSingle NoExtField sigTy -> parens N $ located sigTy p_hsSigType
        DctMulti NoExtField sigTys ->
          parens N $
            sep
              commaDel
              (sitcc . located' p_hsSigType)
              sigTys
  space
  case deriv_clause_strategy of
    Nothing -> do
      breakpoint
      inci derivingWhat
    Just (L _ a) -> case a of
      StockStrategy _ -> do
        txt "stock"
        breakpoint
        inci derivingWhat
      AnyclassStrategy _ -> do
        txt "anyclass"
        breakpoint
        inci derivingWhat
      NewtypeStrategy _ -> do
        txt "newtype"
        breakpoint
        inci derivingWhat
      ViaStrategy (XViaStrategyPs _ sigTy) -> do
        breakpoint
        inci $ do
          derivingWhat
          breakpoint
          txt "via"
          space
          located sigTy p_hsSigType

----------------------------------------------------------------------------
-- Helpers

isInfix :: LexicalFixity -> Bool
isInfix = \case
  Infix -> True
  Prefix -> False

isSingleConstRec :: [LConDecl GhcPs] -> Bool
isSingleConstRec [(L _ ConDeclH98 {..})] =
  case con_args of
    RecCon _ -> True
    _ -> False
isSingleConstRec _ = False

hasHaddocks :: [LConDecl GhcPs] -> Bool
hasHaddocks = any (f . unLoc)
  where
    f ConDeclH98 {..} =
      isJust con_doc || case con_args of
        PrefixCon [] xs ->
          conArgsHaveHaddocks (hsScaledThing <$> xs)
        _ -> False
    f _ = False

conArgsHaveHaddocks :: [LBangType GhcPs] -> Bool
conArgsHaveHaddocks xs =
  let hasDocs = \case
        HsDocTy {} -> True
        _ -> False
   in any (hasDocs . unLoc) xs