Pretty.hs

module Axel.Pretty where

import Axel.Prelude

import Axel.Parse (unhygenisizeIdentifier)
import Axel.Parse.AST
import qualified Axel.Sourcemap as SM
import Axel.Utils.Foldable (mapWithPrev)
import Axel.Utils.Recursion (bottomUpFmap)

import Control.Lens (ala, under)

import Data.Semigroup (Max(Max))
import qualified Data.Text as T
import Data.Text.Lens (unpacked)
import qualified Prettyprinter as P
import Prettyprinter ((<+>))
import qualified Prettyprinter.Render.Text as P

render :: P.Doc a -> Text
render = P.renderStrict . P.layoutSmart (P.LayoutOptions P.Unbounded)

maximumExpressionLength :: Int
maximumExpressionLength = 60

columnWidth :: Text -> Int
columnWidth = ala Max foldMap . map T.length . T.lines

sexp :: Bool -> [P.Doc a] -> P.Doc a
sexp _ [] = mempty
sexp _ [x] = x
sexp canBalance (x:xs) =
  let align cons = x `cons` xs
      oneLiner = align (\y ys -> y <+> P.hsep ys)
      balanced = align (\y ys -> y <+> P.align (P.vsep ys))
      multiLiner = align (\y ys -> P.align (P.vsep (y : ys)))
      shortEnough doc = columnWidth (render doc) <= maximumExpressionLength
   in if | shortEnough oneLiner -> oneLiner
         | canBalance && shortEnough balanced -> balanced
         | otherwise -> multiLiner

privilegedFormToAxelPretty :: [Expression ann] -> P.Doc a
privilegedFormToAxelPretty (Symbol _ "data":name:rest) =
  "data" <+> sexp True (toAxelPretty name : map toAxelPretty rest)
privilegedFormToAxelPretty (Symbol _ "def":name:tySig:rest) =
  "def" <+>
  P.align
    (P.vsep
       [ toAxelPretty name <+> toAxelPretty tySig
       , sexp True (map toAxelPretty rest)
       ])
privilegedFormToAxelPretty (Symbol _ "defmacro":name:rest) =
  "defmacro" <+> sexp True (toAxelPretty name : map toAxelPretty rest)
privilegedFormToAxelPretty (Symbol _ "::":name:rest) =
  "::" <+> sexp True (toAxelPretty name : map toAxelPretty rest)
privilegedFormToAxelPretty (Symbol _ "=":name:args:rest) =
  "=" <+>
  sexp True ((toAxelPretty name <+> toAxelPretty args) : map toAxelPretty rest)
privilegedFormToAxelPretty (Symbol _ "=macro":name:args:rest) =
  "=macro" <+>
  sexp True ((toAxelPretty name <+> toAxelPretty args) : map toAxelPretty rest)
privilegedFormToAxelPretty (Symbol _ "let":bindings:rest) =
  "let" <+> sexp True (toAxelPretty bindings : map toAxelPretty rest)
privilegedFormToAxelPretty (Symbol _ "if":bindings:rest) =
  "if" <+> sexp True (toAxelPretty bindings : map toAxelPretty rest)
privilegedFormToAxelPretty (Symbol _ "case":bindings:rest) =
  "case" <+> sexp True (toAxelPretty bindings : map toAxelPretty rest)
privilegedFormToAxelPretty (Symbol _ "import":name:rest) =
  "import" <+> sexp True (toAxelPretty name : map toAxelPretty rest)
privilegedFormToAxelPretty (Symbol _ "importq":name:rest) =
  "importq" <+> sexp True (toAxelPretty name : map toAxelPretty rest)
privilegedFormToAxelPretty (Symbol _ "importm":name:rest) =
  "importm" <+> sexp True (toAxelPretty name : map toAxelPretty rest)
privilegedFormToAxelPretty xs = sexp True $ map toAxelPretty xs

toAxelPretty :: Expression ann -> P.Doc a
toAxelPretty (LiteralChar _ x) = "#\\" <> P.pretty x
toAxelPretty (LiteralFloat _ x) = P.pretty x
toAxelPretty (LiteralInt _ x) = P.pretty x
toAxelPretty (LiteralString _ x) =
  P.dquotes $ P.pretty (under unpacked showText x)
toAxelPretty (SExpression _ (Symbol _ "applyInfix":xs)) =
  P.braces $ sexp True (map toAxelPretty xs)
toAxelPretty (SExpression _ (Symbol _ "list":xs)) =
  P.brackets $ sexp False (map toAxelPretty xs)
toAxelPretty (SExpression _ [Symbol _ "quote", x]) = "\'" <> toAxelPretty x
toAxelPretty (SExpression _ [Symbol _ "quasiquote", x]) = "`" <> toAxelPretty x
toAxelPretty (SExpression _ [Symbol _ "unquote", x]) = "~" <> toAxelPretty x
toAxelPretty (SExpression _ [Symbol _ "unquoteSplicing", x]) =
  "~@" <> toAxelPretty x
toAxelPretty (Symbol _ x) = P.pretty x
toAxelPretty (SExpression _ xs) = P.parens $ privilegedFormToAxelPretty xs

spaceStatements :: [SM.Expression] -> [P.Doc a]
spaceStatements =
  mapWithPrev $ \prev x ->
    let isImport (SExpression _ (Symbol _ id':_)) =
          id' `elem` ["import", "importq", "importm"]
        isImport _ = False
        needsSpacing = not $ isImport x && maybe False isImport prev
        maybeSpacer =
          if needsSpacing
            then P.line
            else mempty
        unhygenisize =
          bottomUpFmap $ \case
            Symbol ann symbol -> Symbol ann $ unhygenisizeIdentifier symbol
            node -> node
     in toAxelPretty (unhygenisize x) <> maybeSpacer

prettifyProgram :: [SM.Expression] -> Text
prettifyProgram = render . P.vsep . spaceStatements