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

{-# OPTIONS_GHC -Wno-orphans #-}

-- | Utilities for rendering Michelson code in a format compatible with
-- Octez software (e.g @octez-client@)
module Morley.Michelson.Printer.Util
  ( RenderDoc(..)
  , printRenderDoc
  , renderOpsList
  , renderOpsListNoBraces
  , wrapInParens
  , renderDocList

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

import Prelude hiding (group)

import Control.Exception (assert)
import Data.Foldable qualified as Foldable
import Fmt (Buildable(build), FromSimpleDoc(..))
import Fmt.Operators ((<+>), (<//>))
import Fmt.Utils (Doc, SimpleDoc, isEmpty, mkLayoutOptions, renderOneLine)
import Prettyprinter
  (align, braces, enclose, encloseSep, group, hsep, layoutPretty, lbracket, parens, punctuate,
  rbracket, semi, sep, space)

-- | 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 'Doc'. Used to pretty print Michelson
-- code and define '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
forall a. Buildable a => a -> Doc
build

instance RenderDoc Doc where
  renderDoc :: RenderContext -> Doc -> Doc
renderDoc = RenderContext -> Doc -> Doc
addParens

-- | 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
forall ann. Doc ann -> Doc ann
align (Doc -> Doc) -> ([a] -> Doc) -> [a] -> Doc
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Doc -> Doc -> Doc -> [Doc] -> Doc
forall ann. Doc ann -> Doc ann -> Doc ann -> [Doc ann] -> Doc ann
encloseSep Doc
forall ann. Doc ann
lbracket Doc
forall ann. Doc ann
rbracket Doc
", " ([Doc] -> Doc) -> ([a] -> [Doc]) -> [a] -> Doc
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> Doc) -> [a] -> [Doc]
forall a b. (a -> b) -> [a] -> [b]
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)

-- | Convert anything renderable to some text format with line width of 80.
printRenderDoc :: (RenderDoc a, FromSimpleDoc b) => Bool -> a -> b
printRenderDoc :: forall a b. (RenderDoc a, FromSimpleDoc b) => Bool -> a -> b
printRenderDoc Bool
oneLine = SimpleDocStream () -> b
forall ann. SimpleDocStream ann -> b
forall a ann. FromSimpleDoc a => SimpleDocStream ann -> a
fmtSimple (SimpleDocStream () -> b) -> (a -> SimpleDocStream ()) -> a -> b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Bool -> Doc -> SimpleDocStream ()
doRender Bool
oneLine (Doc -> SimpleDocStream ())
-> (a -> Doc) -> a -> SimpleDocStream ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. RenderContext -> a -> Doc
forall a. RenderDoc a => RenderContext -> a -> Doc
renderDoc RenderContext
doesntNeedParens

-- | Render a comma-separated list of items in braces
renderOpsList :: (RenderDoc op, Foldable f) => Bool -> f op -> Doc
renderOpsList :: forall op (f :: * -> *).
(RenderDoc op, Foldable f) =>
Bool -> f op -> Doc
renderOpsList Bool
oneLine f op
ops =
  Doc -> Doc
forall ann. Doc ann -> Doc ann
braces (Doc -> Doc) -> Doc -> Doc
forall a b. (a -> b) -> a -> b
$ Doc -> Doc -> Doc -> Doc
forall ann. Doc ann -> Doc ann -> Doc ann -> Doc ann
enclose Doc
forall ann. Doc ann
space Doc
forall ann. Doc ann
space (Doc -> Doc) -> Doc -> Doc
forall a b. (a -> b) -> a -> b
$ Bool -> f op -> Doc
forall op (f :: * -> *).
(RenderDoc op, Foldable f) =>
Bool -> f op -> Doc
renderOpsListNoBraces Bool
oneLine f op
ops

-- | Render a semi-colon-separated list of items without braces
renderOpsListNoBraces :: (RenderDoc op, Foldable f) => Bool -> f op -> Doc
renderOpsListNoBraces :: forall op (f :: * -> *).
(RenderDoc op, Foldable f) =>
Bool -> f op -> Doc
renderOpsListNoBraces Bool
oneLine =
  Doc -> Doc
forall ann. Doc ann -> Doc ann
align (Doc -> Doc) -> (f op -> Doc) -> f op -> Doc
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (if Bool
oneLine then [Doc] -> Doc
forall ann. [Doc ann] -> Doc ann
hsep else [Doc] -> Doc
forall ann. [Doc ann] -> Doc ann
sep) ([Doc] -> Doc) -> (f op -> [Doc]) -> f op -> Doc
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Doc -> [Doc] -> [Doc]
forall ann. Doc ann -> [Doc ann] -> [Doc ann]
punctuate Doc
forall ann. Doc ann
semi ([Doc] -> [Doc]) -> (f op -> [Doc]) -> f op -> [Doc]
forall b c a. (b -> c) -> (a -> b) -> a -> c
.
    (op -> Doc) -> [op] -> [Doc]
forall a b. (a -> b) -> [a] -> [b]
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]) -> (f op -> [op]) -> f 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 ([op] -> [op]) -> (f op -> [op]) -> f op -> [op]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. f op -> [op]
forall a. f a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
Foldable.toList

-- | 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 i a.
(Integral i, Container a,
 DefaultToInt (IsIntSubType Length i) i) =>
a -> i
length ((Doc -> Bool) -> [Doc] -> [Doc]
forall a. (a -> Bool) -> [a] -> [a]
filter (Bool -> Bool
forall a. Boolean a => a -> a
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
forall b.
(Element (NonEmpty Doc) -> b -> b) -> b -> NonEmpty Doc -> b
foldr Doc -> Doc -> Doc
Element (NonEmpty Doc) -> Doc -> Doc
(<+>) Doc
forall a. Monoid a => a
mempty NonEmpty Doc
ds

instance {-# overlappable #-} RenderDoc a => Buildable a where
  build :: a -> Doc
build = Doc -> Doc
forall ann. Doc ann -> Doc ann
align (Doc -> Doc) -> (a -> Doc) -> a -> Doc
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Doc -> Doc
forall ann. Doc ann -> Doc ann
group (Doc -> Doc) -> (a -> Doc) -> a -> Doc
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 -> SimpleDocStream ()
doRender Bool
oneLine = if Bool
oneLine then Doc -> SimpleDocStream ()
renderOneLine else LayoutOptions -> Doc -> SimpleDocStream ()
forall ann. LayoutOptions -> Doc ann -> SimpleDocStream ann
layoutPretty (LayoutOptions -> Doc -> SimpleDocStream ())
-> LayoutOptions -> Doc -> SimpleDocStream ()
forall a b. (a -> b) -> a -> b
$ Double -> Int -> LayoutOptions
mkLayoutOptions Double
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
forall ann. Doc ann -> Doc ann
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
forall a. Boolean a => a -> a
not Bool
pn)