Make printer more explicit about one-line vs multi-line strings

src/lib/Name.hs

@@ -1742,7 +1742,7 @@ instance RenameE e => RenameE (NonEmptyListE e) where
   renameE env (NonEmptyListE xs) = NonEmptyListE $ fmap (renameE env) xs
 
 instance (PrettyB b, PrettyE e) => Pretty (Abs b e n) where
-  pr (Abs b body) = hcat [pr b, indent (pr body)]
+  pr (Abs b body) = pr b <+> pr body
 
 instance Pretty a => Pretty (LiftE a n) where
   pr (LiftE x) = pr x
@@ -2665,7 +2665,7 @@ instance HoistableB b => HoistableB (RNest b) where
 
 instance (forall n. Pretty (v n)) => Pretty (SubstFrag v i i' o) where
   pr (UnsafeMakeSubst m) =
-    vcat [ hcat [pr v, " @> ", pr x] | (v, SubstItem _ x) <- R.toList m ]
+    hcat [ hcat [pr v, " @> ", pr x] | (v, SubstItem _ x) <- R.toList m ]
 
 instance (Generic (b UnsafeS UnsafeS)) => Generic (Nest b n l) where
   type Rep (Nest b n l) = Rep [b UnsafeS UnsafeS]

src/lib/PPrint.hs

@@ -4,100 +4,95 @@
 -- license that can be found in the LICENSE file or at
 -- https://developers.google.com/open-source/licenses/bsd
 
-module PPrint (Pretty (..), Doc (..), indent, hcat, hlist, vcat, pprint, app) where
+{-# LANGUAGE NoFieldSelectors #-}
+
+module PPrint (
+  Pretty (..), indent, emitLine, hcat, hlist, pprint, app,
+  (<+>), BSBuilder, forceOneLine) where
 
 import Data.Int
 import Data.Word
 import Data.List (intersperse)
 import qualified Data.ByteString as BS
+import qualified Data.ByteString.Builder as BS
 import Data.String
-import Control.Monad.Reader
 import Control.Monad.State.Strict
 
 pprint :: Pretty a => a -> BString
-pprint x = printDoc $ pr x
+pprint x = runPrinter $ prLines x
 {-# SCC pprint #-}
 
 -- === printing doc ===
 
 type BString = BS.ByteString
+type Indent = BS.Builder
+type BSBuilder = BS.Builder
 
-type Indent = BString
-
-newtype PrinterM a = PrinterM { runPrinterM :: ReaderT Indent (State [(Indent, BString)]) a }
+data PrinterState = PrinterState {indent :: Indent, curString  :: BS.Builder }
+newtype PrinterM a = PrinterM { inner :: State PrinterState a }
         deriving (Functor, Applicative, Monad)
 
-runPrinter :: PrinterM a -> BString
-runPrinter cont = do
-  let indentedLines = reverse $ execState (runReaderT (runPrinterM cont) "") []
-  BS.concat [indent <> s <> "\n"| (indent, s) <- indentedLines]
-
-printDoc :: Doc -> BString
-printDoc d = runPrinter $ printDocM d
-
-increaseIndent :: PrinterM a -> PrinterM a
-increaseIndent cont = PrinterM $ local (<> "  ") $ runPrinterM cont
-
-printDocM :: Doc -> PrinterM ()
-printDocM = \case
-  DocLine s -> do
-    curIndent <- PrinterM ask
-    PrinterM $ modify \indentedLines -> (curIndent, s) : indentedLines
-  DocIndent d -> increaseIndent $ printDocM d
-  DocItems ds -> mapM_ printDocM ds
-
--- === constructing doc ===
-
+-- Instances should define either `pr` (if they're expected to be one-liners
+-- most of the time) or `prLines`.
 class Pretty a where
-  pr :: a -> Doc
+  pr :: a -> BSBuilder
+  pr x = forceOneLine $ prLines x
 
-  prList :: [a] -> Doc
-  prList xs = hlist "[,]" $ map pr xs
-
-data Doc =
-   DocLine BString
- | DocItems  [Doc]
- | DocIndent Doc
-   deriving (Show)
+  prLines :: a -> PrinterM ()
+  prLines x = emitLine $ pr x
 
-vcat :: [Doc] -> Doc
-vcat = DocItems
+  prList :: [a] -> BS.Builder
+  prList xs = hlist "[,]" $ map pr xs
 
-hlist :: String -> [Doc] -> Doc
+runPrinter :: PrinterM a -> BString
+runPrinter cont = BS.toStrict $ BS.toLazyByteString $ (.curString) $
+  execState cont.inner $ PrinterState mempty mempty
+
+-- This is a fallback and we shouldn't see its output much
+forceOneLine :: PrinterM () -> BSBuilder
+forceOneLine x = "\n{" <> BS.byteString (runPrinter x) <> "}\n"
+
+indent :: PrinterM a -> PrinterM a
+indent cont = PrinterM do
+  prev <- gets (.indent)
+  modify \s -> s {indent = prev <> "  "}
+  ans <- cont.inner
+  modify \s -> s {indent = prev}
+  return ans
+
+emitLine :: BS.Builder -> PrinterM ()
+emitLine b = PrinterM do
+  s <- get
+  put $ s {curString = s.curString <> "\n" <> s.indent <> b}
+
+hlist :: String -> [BS.Builder] -> BS.Builder
 hlist [l,sep,r] xs = hcat [pr l, hcat (intersperse (pr sep) xs), pr r]
-
 hlist _ _ = error "expected left bracket, separator, right bracket"
 
-hcat :: [Doc] -> Doc
-hcat docs = rec "" docs
- where
-  rec :: BString -> [Doc] -> Doc
-  rec s = \case
-    [] -> DocLine s
-    d:ds -> case d of
-      DocLine s' -> rec (s <> s') ds
-      _ -> vcat [DocLine s, d, hcat ds]
-
-indent :: Doc -> Doc
-indent = DocIndent
+hcat :: [BS.Builder] -> BS.Builder
+hcat = mconcat
 
-app :: Doc -> [Doc] -> Doc
+app :: BS.Builder -> [BS.Builder] -> BS.Builder
 app f xs = hcat [f, hlist "(,)" xs]
 
+infixr 6 <+>
+(<+>) :: BS.Builder -> BS.Builder -> BS.Builder
+(<+>) x y = hcat [x, " ", y]
+
 -- === instances ===
 
-instance IsString Doc where
-  fromString s = DocLine $ fromString s
+instance IsString (PrinterM ()) where
+  fromString s = emitLine $ fromString s
 
 instance Pretty Char where
-  pr c = DocLine $ fromString [c]
-  prList s = DocLine $ fromString s
+  pr c = fromString [c]
+  prList s = fromString s
 
 instance Pretty a => Pretty [a] where
   pr xs = prList xs
 
 instance Pretty BString where
-  pr s = DocLine s
+  pr s = BS.byteString s
 
 instance (Pretty a, Pretty b) => Pretty (a, b) where
   pr (x, y) = hcat ["(", pr x, ", ", pr y, ")"]

src/lib/Simplify.hs

@@ -31,8 +31,8 @@ exprAsNullaryFun :: CExpr VoidS -> CoreLamExpr VoidS
 exprAsNullaryFun expr = CoreLamExpr piType (Abs Empty expr)
   where piType = CorePiType ExplicitApp [] (Abs Empty (getCType expr))
 
-simplifyTopFun :: Monad m => CoreLamExpr VoidS -> m (LamExpr VoidS)
-simplifyTopFun f = liftSimplifyM $ simplifyLam f
+simplifyTopFun :: Monad m => CoreLamExpr VoidS -> m TopLamExpr
+simplifyTopFun f = liftSimplifyM $ TopLamExpr <$> simplifyLam f
 
 -- === Simplification monad ===
 

src/lib/TopLevel.hs

@@ -108,10 +108,10 @@ execUDecl decl = do
   llvmContext <- TopperM $ asks topperLLVMContext
   llvmFun <- toLLVMEntryFun tempFunName simpFun
   logPass LLVMPass llvmFun
-  liftIO do
-    compileLLVM llvmContext llvmFun
-    f <- getFunctionPtr llvmContext tempFunName
-    callEntryFun f []
+  -- liftIO do
+  --   compileLLVM llvmContext llvmFun
+  --   f <- getFunctionPtr llvmContext tempFunName
+  --   callEntryFun f []
   return ()
 
 execCDecl :: CTopDecl -> TopperM ()

src/lib/Types/Complicated.hs

@@ -19,7 +19,6 @@ import GHC.Generics (Generic (..))
 import Data.Store (Store (..))
 
 import Name
-import Util (Tree (..))
 import PPrint
 
 import Types.Source (HasSourceName (..))

src/lib/Types/Primitives.hs

@@ -339,22 +339,22 @@ instance Pretty Direction where
     Fwd -> "fwd"
     Rev -> "rev"
 
-printDouble :: Double -> Doc
-printDouble x = pr (double2Float x)
+printDouble :: Double -> String
+printDouble x = printFloat $ double2Float x
 
-printFloat :: Float -> Doc
-printFloat x = pr $ reverse $ dropWhile (=='0') $ reverse $
+printFloat :: Float -> String
+printFloat x = fromString $ reverse $ dropWhile (=='0') $ reverse $
   showFFloat (Just 6) x ""
 
 instance Pretty LitVal where
   pr = \case
     Int64Lit   x -> pr x
     Int32Lit   x -> pr x
-    Float64Lit x -> printDouble x
-    Float32Lit x -> printFloat  x
-    Word8Lit   x -> pr $ show $ toEnum @Char $ fromIntegral x
-    Word32Lit  x -> pr $ "0x" ++ showHex x ""
-    Word64Lit  x -> pr $ "0x" ++ showHex x ""
+    Float64Lit x -> fromString $ printDouble x
+    Float32Lit x -> fromString $ printFloat  x
+    Word8Lit   x -> fromString $ show $ toEnum @Char $ fromIntegral x
+    Word32Lit  x -> fromString $ "0x" ++ showHex x ""
+    Word64Lit  x -> fromString $ "0x" ++ showHex x ""
     PtrLit ty (PtrLitVal x) -> app "Ptr" [pr ty, pr (show x)]
     PtrLit _ NullPtr -> "NullPtr"
     PtrLit _ (PtrSnapshot _) -> "<ptr snapshot>"

src/lib/Types/Simple.hs

@@ -1,4 +1,3 @@
-
 -- Copyright 2022 Google LLC
 --
 -- Use of this source code is governed by a BSD-style
@@ -10,6 +9,7 @@
 
 module Types.Simple (module Types.Simple) where
 
+import Control.Monad
 import Data.Word
 import Data.Foldable (toList)
 import Data.Hashable
@@ -20,14 +20,15 @@ import GHC.Generics (Generic (..))
 import Data.Store (Store (..))
 
 import Name
-import Util (Tree (..))
 import PPrint
 
 import Types.Primitives
 import Types.Source (HasSourceName (..))
 
 -- === SimpIR ===
 
+data TopLamExpr = TopLamExpr (LamExpr VoidS)
+
 data Expr (n::S) =
    Block  (Type n) (Block n)
  | TopApp (Type n) TopName [Atom n]
@@ -66,6 +67,11 @@ type PiType = Abs (Nest Binder) Type :: E
 
 -- === type classes ===
 
+instance Pretty TopLamExpr where
+  prLines (TopLamExpr (Abs bs body)) = do
+    emitLine (pr bs <> ".")
+    indent $ prLines body
+
 instance GenericE Expr where
   type RepE Expr = EitherE6
   {- Block  -} (Type `PairE` Block)
@@ -85,14 +91,23 @@ instance GenericE Expr where
 
 instance Pretty (Expr n) where
   pr = \case
-    Block _ (Abs decls result) ->
-      vcat (nestToList' pr decls ++ [pr result])
+    Block _ (Abs decls result) -> undefined
     TopApp _ _ _ -> undefined
     Case   _ _ _ -> undefined
     For    _ _ -> undefined
     While  _ -> undefined
     PrimOp _ op -> pr op
 
+  prLines = \case
+    Block _ (Abs decls result) -> do
+      forM_ (nestToList' pr decls) emitLine
+      emitLine $ pr result
+    TopApp _ _ _ -> undefined
+    Case   _ _ _ -> undefined
+    For    _ _ -> undefined
+    While  _ -> undefined
+    PrimOp _ op -> prLines op
+
 instance SinkableE      Expr
 instance HoistableE     Expr
 instance RenameE        Expr

src/lib/Types/Source.hs

@@ -989,13 +989,15 @@ instance Pretty SourceBlock where
   pr block = pr $ sbText block
 
 instance Pretty Output where
-  pr = \case
-    TextOut s -> pr s
+  prLines = \case
+    TextOut s -> prLines s
     HtmlOut _ -> "<html output>"
     SourceInfo _ -> "<source info>"
-    PassResult name s -> vcat [hcat [" === ", pr name, " ==="], pr s]
-    MiscLog s -> pr s
-    Error e -> pr e
+    PassResult name s -> do
+      emitLine $ hcat [" === ", pr name, " ==="]
+      prLines s
+    MiscLog s -> prLines s
+    Error e -> prLines e
 
 instance Pretty PassName where
   pr x = pr $ show x

src/lib/Util.hs

@@ -11,6 +11,7 @@ module Util where
 import Prelude
 import qualified Data.Set as Set
 import qualified Data.Map.Strict as M
+import Data.ByteString.Internal (w2c)
 import Control.Applicative
 import Control.Monad
 import Control.Monad.Reader
@@ -328,28 +329,11 @@ tryUnsnoc (ReversedList []) = Nothing
 tryUnsnoc (ReversedList (x:xs)) = Just (ReversedList xs, x)
 {-# INLINE tryUnsnoc #-}
 
--- === generic tree ===
-
-data Tree a = Leaf a | Branch [Tree a]
-     deriving (Show, Eq, Ord, Generic, Functor, Foldable, Traversable)
-instance Store a => Store (Tree a)
-instance Hashable a => Hashable (Tree a)
-
-zipTrees :: Tree a -> Tree b -> Tree (a, b)
-zipTrees (Leaf x) (Leaf y) = Leaf (x, y)
-zipTrees (Branch xs) (Branch ys) | length xs == length ys = Branch $ zipWith zipTrees xs ys
-zipTrees _ _ = error "zip error"
-
-instance Pretty a => Pretty (Tree a) where
-  pr = \case
-    Leaf x -> pr x
-    Branch xs -> pr xs
+-- === bytestring <-> string conversion ===
 
 readFileText :: MonadIO m => FilePath -> m BString
 readFileText fname = liftIO $ BS.readFile fname
 
--- === bytestring <-> string conversion ===
-
 type BString = BS.ByteString
 
 showbs :: Show a => a -> BString
@@ -359,7 +343,7 @@ errorbs :: HasCallStack => BString -> a
 errorbs s = error $ bs2str s
 
 bs2str :: BString -> String
-bs2str = undefined
+bs2str s = map w2c $ BS.unpack s
 
 str2bs :: String -> BString
 str2bs = fromString