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

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

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

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, vcat, (<+>))

-- | Environment carried during recursive rendering.
newtype RenderContext = RenderContext
  { RenderContext -> Bool
_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 _ = Bool
True

-- | 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 :: Bool -> Doc -> Text
printDoc oneLine :: Bool
oneLine = SimpleDoc -> Text
displayT (SimpleDoc -> Text) -> (Doc -> SimpleDoc) -> Doc -> Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Bool -> Doc -> SimpleDoc
doRender Bool
oneLine

-- | Convert 'Doc' to 'Builder' in the same maner as 'printDoc'.
printDocB :: Bool -> Doc -> Builder
printDocB :: Bool -> Doc -> Builder
printDocB oneLine :: Bool
oneLine = SimpleDoc -> Builder
displayB (SimpleDoc -> Builder) -> (Doc -> SimpleDoc) -> Doc -> Builder
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Bool -> Doc -> SimpleDoc
doRender Bool
oneLine

-- | Convert 'Doc' to 'String' in the same maner as 'printDoc'.
printDocS :: Bool -> Doc -> String
printDocS :: Bool -> Doc -> String
printDocS oneLine :: Bool
oneLine = Text -> String
forall a. ToString a => a -> String
toString (Text -> String) -> (Doc -> Text) -> Doc -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Bool -> Doc -> Text
printDoc Bool
oneLine

-- | Generic way to render the different op types that get passed
-- to a contract.
renderOps :: (RenderDoc op) => Bool -> NonEmpty op -> Doc
renderOps :: Bool -> NonEmpty op -> Doc
renderOps oneLine :: Bool
oneLine = Bool -> [op] -> Doc
forall op. RenderDoc op => Bool -> [op] -> Doc
renderOpsList Bool
oneLine ([op] -> Doc) -> (NonEmpty op -> [op]) -> NonEmpty op -> Doc
forall b c a. (b -> c) -> (a -> b) -> a -> c
. NonEmpty op -> [op]
forall t. Container t => t -> [Element t]
toList

spacecat :: NonEmpty Doc -> Doc
spacecat :: NonEmpty Doc -> Doc
spacecat = (Element (NonEmpty Doc) -> Doc -> Doc)
-> Doc -> NonEmpty Doc -> Doc
forall t b. Container t => (Element t -> b -> b) -> b -> t -> b
foldr Element (NonEmpty Doc) -> Doc -> Doc
Doc -> Doc -> Doc
(<+>) Doc
forall a. Monoid a => a
mempty

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

renderOpsListNoBraces :: RenderDoc op => Bool -> [op] -> Doc
renderOpsListNoBraces :: Bool -> [op] -> Doc
renderOpsListNoBraces oneLine :: Bool
oneLine ops :: [op]
ops =
  [Doc] -> Doc
cat' ([Doc] -> Doc) -> [Doc] -> Doc
forall a b. (a -> b) -> a -> b
$ Doc -> [Doc] -> [Doc]
punctuate Doc
semi ([Doc] -> [Doc]) -> [Doc] -> [Doc]
forall a b. (a -> b) -> a -> b
$
    RenderContext -> op -> Doc
forall a. RenderDoc a => RenderContext -> a -> Doc
renderDoc RenderContext
doesntNeedParens (op -> Doc) -> [op] -> [Doc]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (op -> Bool) -> [op] -> [op]
forall a. (a -> Bool) -> [a] -> [a]
filter op -> Bool
forall a. RenderDoc a => a -> Bool
isRenderable [op]
ops
  where
    cat' :: [Doc] -> Doc
cat' = if Bool
oneLine then Doc -> (NonEmpty Doc -> Doc) -> Maybe (NonEmpty Doc) -> Doc
forall b a. b -> (a -> b) -> Maybe a -> b
maybe "" NonEmpty Doc -> Doc
spacecat (Maybe (NonEmpty Doc) -> Doc)
-> ([Doc] -> Maybe (NonEmpty Doc)) -> [Doc] -> Doc
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Doc] -> Maybe (NonEmpty Doc)
forall a. [a] -> Maybe (NonEmpty a)
nonEmpty else Doc -> Doc
align (Doc -> Doc) -> ([Doc] -> Doc) -> [Doc] -> Doc
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Doc] -> Doc
vcat

-- | Create a specific number of spaces.
spaces :: Int -> Doc
spaces :: Int -> Doc
spaces x :: Int
x = [Doc] -> Doc
hcat ([Doc] -> Doc) -> [Doc] -> Doc
forall a b. (a -> b) -> a -> b
$ Int -> Doc -> [Doc]
forall a. Int -> a -> [a]
replicate Int
x Doc
space

-- | Wrap documents in parentheses if there are two or more in the list.
wrapInParens :: RenderContext -> NonEmpty Doc -> Doc
wrapInParens :: RenderContext -> NonEmpty Doc -> Doc
wrapInParens (RenderContext pn :: Bool
pn) ds :: NonEmpty Doc
ds =
  let moreThanOne :: Bool
moreThanOne = [Doc] -> Int
forall t. Container t => t -> Int
length ((Doc -> Bool) -> [Doc] -> [Doc]
forall a. (a -> Bool) -> [a] -> [a]
filter (Bool -> Bool
not (Bool -> Bool) -> (Doc -> Bool) -> Doc -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Doc -> Bool
isEmpty) (NonEmpty Doc -> [Element (NonEmpty Doc)]
forall t. Container t => t -> [Element t]
toList NonEmpty Doc
ds)) Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> 1
  in  RenderContext -> Doc -> Doc
addParens (Bool -> RenderContext
RenderContext (Bool
pn Bool -> Bool -> Bool
&& Bool
moreThanOne)) (Doc -> Doc) -> Doc -> Doc
forall a b. (a -> b) -> a -> b
$
        (Element (NonEmpty Doc) -> Doc -> Doc)
-> Doc -> NonEmpty Doc -> Doc
forall t b. Container t => (Element t -> b -> b) -> b -> t -> b
foldr Element (NonEmpty Doc) -> Doc -> Doc
Doc -> Doc -> Doc
(<+>) Doc
forall a. Monoid a => a
mempty NonEmpty Doc
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 :: a -> Builder
buildRenderDoc = Bool -> Doc -> Builder
printDocB Bool
True (Doc -> Builder) -> (a -> Doc) -> a -> Builder
forall b c a. (b -> c) -> (a -> b) -> a -> c
. RenderContext -> a -> Doc
forall a. RenderDoc a => RenderContext -> a -> Doc
renderDoc RenderContext
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 :: Bool -> Doc -> SimpleDoc
doRender oneLine :: Bool
oneLine = if Bool
oneLine then Doc -> SimpleDoc
renderOneLine else Float -> Int -> Doc -> SimpleDoc
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
needsParens = Bool -> RenderContext
RenderContext Bool
True
doesntNeedParens :: RenderContext
doesntNeedParens = Bool -> RenderContext
RenderContext Bool
False

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

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