-- SPDX-FileCopyrightText: 2021 Oxhead Alpha
-- SPDX-License-Identifier: LicenseRef-MIT-OA

-- | Utilities for rendering Michelson code in a format compatible with
-- Octez software (e.g @octez-client@)
module Morley.Michelson.Printer.Util
  ( RenderDoc(..)
  , Prettier(..)
  , printDoc
  , printDocB
  , printDocS
  , renderOps
  , renderOpsList
  , renderOpsListNoBraces
  , renderAnyBuildable
  , wrapInParens
  , buildRenderDoc
  , buildRenderDocExtended
  , renderDocList

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

import Prelude hiding (group)

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

-- | 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 sufficient 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

instance RenderDoc Text where
  renderDoc :: RenderContext -> Text -> Doc
renderDoc RenderContext
_ = Text -> Doc
textStrict

-- | Renders a list of 'RenderDoc' elements surrounded with square brackets,
-- separated by a comma and a space.
renderDocList :: RenderDoc a => RenderContext -> [a] -> Doc
renderDocList :: forall a. RenderDoc a => 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 :: forall a. Buildable a => 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
  deriving stock (forall a b. (a -> b) -> Prettier a -> Prettier b)
-> (forall a b. a -> Prettier b -> Prettier a) -> Functor Prettier
forall a b. a -> Prettier b -> Prettier a
forall a b. (a -> b) -> Prettier a -> Prettier b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
<$ :: forall a b. a -> Prettier b -> Prettier a
$c<$ :: forall a b. a -> Prettier b -> Prettier a
fmap :: forall a b. (a -> b) -> Prettier a -> Prettier b
$cfmap :: forall a b. (a -> b) -> Prettier a -> Prettier b
Functor

-- | 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 manner 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 manner 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 :: forall op. RenderDoc op => 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

-- | Render a comma-separated list of items in braces
renderOpsList :: (RenderDoc op) => Bool -> [op] -> Doc
renderOpsList :: forall op. RenderDoc op => 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

-- | Render a semi-colon-separated list of items without braces
renderOpsListNoBraces :: RenderDoc op => Bool -> [op] -> Doc
renderOpsListNoBraces :: forall op. RenderDoc op => Bool -> [op] -> Doc
renderOpsListNoBraces Bool
oneLine =
  Doc -> Doc
align (Doc -> Doc) -> ([op] -> Doc) -> [op] -> Doc
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (if Bool
oneLine then [Doc] -> Doc
hsep else [Doc] -> Doc
sep) ([Doc] -> Doc) -> ([op] -> [Doc]) -> [op] -> Doc
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Doc -> [Doc] -> [Doc]
punctuate Doc
semi ([Doc] -> [Doc]) -> ([op] -> [Doc]) -> [op] -> [Doc]
forall b c a. (b -> c) -> (a -> b) -> a -> c
.
    (op -> Doc) -> [op] -> [Doc]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (RenderContext -> op -> Doc
forall a. RenderDoc a => RenderContext -> a -> Doc
renderDoc RenderContext
doesntNeedParens) ([op] -> [Doc]) -> ([op] -> [op]) -> [op] -> [Doc]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (op -> Bool) -> [op] -> [op]
forall a. (a -> Bool) -> [a] -> [a]
filter op -> Bool
forall a. RenderDoc a => a -> Bool
isRenderable

-- | 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 :: forall a. RenderDoc a => 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 :: forall a. RenderDoc a => 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.
-}

-- | Constructors for 'RenderContext'
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

-- | Add parentheses if needed, multiline if necessary.
addParensMultiline :: RenderContext -> Doc -> Doc
addParensMultiline :: RenderContext -> Doc -> Doc
addParensMultiline RenderContext
pn Doc
doc = case RenderContext
pn of
  RenderContext Bool
True -> Doc
"(" Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Doc
doc Doc -> Doc -> Doc
<//> Doc
")"
  RenderContext Bool
False -> Doc
doc

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