Util.hs

-- SPDX-FileCopyrightText: 2020 Tocqueville Group
--
-- SPDX-License-Identifier: LicenseRef-MIT-TQ

module Michelson.Printer.Util
  ( RenderDoc(..)
  , Prettier(..)
  , printDoc
  , printDocB
  , printDocS
  , renderOps
  , renderOpsList
  , renderOpsListNoBraces
  , renderAnyBuildable
  , spaces
  , wrapInParens
  , buildRenderDoc
  , buildRenderDocExtended
  , renderDocList

    -- * Smart parentheses
  , RenderContext
  , needsParens
  , doesntNeedParens
  , addParens
  , assertParensNotNeeded
  ) where

import Fmt (Buildable, pretty)
import Control.Exception (assert)
import qualified Data.Text.Lazy as LT
import Data.Text.Lazy.Builder (Builder)
import Text.PrettyPrint.Leijen.Text
  (Doc, SimpleDoc, align, braces, displayB, displayT, enclose, hcat, isEmpty, parens, punctuate,
  renderOneLine, renderPretty, semi, space, text, vcat, (<+>), lbracket, rbracket, encloseSep)

-- | Environment carried during recursive rendering.
newtype RenderContext = RenderContext
  { _rcWillNeedParens :: Bool
    -- ^ Whether the current expression is going to be used as part of
    -- top-level expression or in a similar context.
    -- When set to 'True', you may need to wrap your rendered expression into
    -- parentheses.
  }

-- | Generalize converting a type into a
-- Text.PrettyPrint.Leijen.Text.Doc. Used to pretty print Michelson code
-- and define Fmt.Buildable instances.
class RenderDoc a where
  renderDoc :: RenderContext -> a -> Doc

  -- | Whether a value can be represented in Michelson code.
  -- Normally either all values of some type are renderable or not renderable.
  -- However, in case of instructions we have extra instructions which should
  -- not be rendered.
  -- Note: it's not suficcient to just return 'mempty' for such instructions,
  -- because sometimes we want to print lists of instructions and we need to
  -- ignore them complete (to avoid putting redundant separators).
  isRenderable :: a -> Bool
  isRenderable _ = True

-- | Renders a list of 'RenderDoc' elements surrounded with square brackets,
-- separated by a comma and a space.
renderDocList :: RenderDoc a => RenderContext -> [a] -> Doc
renderDocList context = encloseSep lbracket rbracket ", " . fmap (renderDoc context)

renderAnyBuildable :: Buildable a => a -> Doc
renderAnyBuildable = text . pretty

-- | A new type that can wrap values so that the RenderDoc
-- instances of the combined value can have a different
-- behavior for the pretty printer.
newtype Prettier a = Prettier a

-- | Convert 'Doc' to 'Text' with a line width of 80.
printDoc :: Bool -> Doc -> LT.Text
printDoc oneLine = displayT . doRender oneLine

-- | Convert 'Doc' to 'Builder' in the same maner as 'printDoc'.
printDocB :: Bool -> Doc -> Builder
printDocB oneLine = displayB . doRender oneLine

-- | Convert 'Doc' to 'String' in the same maner as 'printDoc'.
printDocS :: Bool -> Doc -> String
printDocS oneLine = toString . printDoc oneLine

-- | Generic way to render the different op types that get passed
-- to a contract.
renderOps :: (RenderDoc op) => Bool -> NonEmpty op -> Doc
renderOps oneLine = renderOpsList oneLine . toList

spacecat :: NonEmpty Doc -> Doc
spacecat = foldr (<+>) mempty

renderOpsList :: (RenderDoc op) => Bool -> [op] -> Doc
renderOpsList oneLine ops =
  braces $ enclose space space $ renderOpsListNoBraces oneLine ops

renderOpsListNoBraces :: RenderDoc op => Bool -> [op] -> Doc
renderOpsListNoBraces oneLine ops =
  cat' $ punctuate semi $
    renderDoc doesntNeedParens <$> filter isRenderable ops
  where
    cat' = if oneLine then maybe "" spacecat . nonEmpty else align . vcat

-- | Create a specific number of spaces.
spaces :: Int -> Doc
spaces x = hcat $ replicate x space

-- | Wrap documents in parentheses if there are two or more in the list.
wrapInParens :: RenderContext -> NonEmpty Doc -> Doc
wrapInParens (RenderContext pn) ds =
  let moreThanOne = length (filter (not . isEmpty) (toList ds)) > 1
  in  addParens (RenderContext (pn && moreThanOne)) $
        foldr (<+>) mempty ds

-- | Turn something that is instance of `RenderDoc` into a `Builder`.
-- It's formatted the same way as `printDoc` formats docs.
buildRenderDoc :: RenderDoc a => a -> Builder
buildRenderDoc = printDocB True . renderDoc doesntNeedParens

-- | Works as 'buildRenderDoc' above, but doesn't force the doc to be printed in one line
buildRenderDocExtended :: RenderDoc a => a -> Builder
buildRenderDocExtended = printDocB False . renderDoc doesntNeedParens


-- | Here using a page width of 80 and a ribbon width of 1.0
-- https://hackage.haskell.org/package/wl-pprint-1.2.1/docs/Text-PrettyPrint-Leijen.html
doRender :: Bool -> Doc -> SimpleDoc
doRender oneLine = if oneLine then renderOneLine else renderPretty 1.0 80

-- Smart parentheses
----------------------------------------------------------------------------

{- Motivation:

Some expressions may need to be wrapped into parentheses, but only if they
are part of other expression, and are not already wrapped into braces or
brackets.
-}

-- | 'ParensNeeded' constant.
needsParens, doesntNeedParens :: RenderContext
needsParens = RenderContext True
doesntNeedParens = RenderContext False

-- | Add parentheses if needed.
addParens :: RenderContext -> Doc -> Doc
addParens = \case
  RenderContext True -> parens
  RenderContext False -> id

-- | Ensure parentheses are not required, for case when you cannot
-- sensibly wrap your expression into them.
assertParensNotNeeded :: RenderContext -> a -> a
assertParensNotNeeded (RenderContext pn) = assert (not pn)