-- 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
  { 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 a
_ = Bool
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 :: RenderContext -> [a] -> Doc
renderDocList RenderContext
context = Doc -> Doc -> Doc -> [Doc] -> Doc
encloseSep Doc
lbracket Doc
rbracket Doc
", " ([Doc] -> Doc) -> ([a] -> [Doc]) -> [a] -> Doc
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> Doc) -> [a] -> [Doc]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (RenderContext -> a -> Doc
forall a. RenderDoc a => RenderContext -> a -> Doc
renderDoc RenderContext
context)

renderAnyBuildable :: Buildable a => a -> Doc
renderAnyBuildable :: a -> Doc
renderAnyBuildable = Text -> Doc
text (Text -> Doc) -> (a -> Text) -> a -> Doc
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Text
forall a b. (Buildable a, FromBuilder b) => a -> b
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 :: Bool -> Doc -> Text
printDoc 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 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 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 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 Bool
oneLine [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 Bool
oneLine [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 Doc
"" 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 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 Bool
pn) 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
> Int
1
  in  RenderContext -> Doc -> Doc
addParens (Bool -> RenderContext
RenderContext (Bool
pn Bool -> Bool -> Bool
forall a. Boolean a => a -> a -> a
&& 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

-- | Works as 'buildRenderDoc' above, but doesn't force the doc to be printed in one line
buildRenderDocExtended :: RenderDoc a => a -> Builder
buildRenderDocExtended :: a -> Builder
buildRenderDocExtended = Bool -> Doc -> Builder
printDocB Bool
False (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 Bool
oneLine = if Bool
oneLine then Doc -> SimpleDoc
renderOneLine else Float -> Int -> Doc -> SimpleDoc
renderPretty Float
1.0 Int
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 Bool
True -> Doc -> Doc
parens
  RenderContext Bool
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 Bool
pn) = Bool -> a -> a
forall a. (?callStack::CallStack) => Bool -> a -> a
assert (Bool -> Bool
not Bool
pn)