--------------------------------------------------------------------------------
{-# LANGUAGE BangPatterns #-}
module Patat.PrettyPrint.Matrix
    ( Matrix
    , Cell (..)
    , emptyCell
    , docToMatrix
    , hPutMatrix
    ) where


--------------------------------------------------------------------------------
import           Control.Monad              (unless, when)
import           Data.Char.WCWidth.Extended (wcwidth)
import qualified Data.Vector                as V
import qualified Data.Vector.Mutable        as VM
import           Patat.PrettyPrint.Internal hiding (null)
import           Patat.Size                 (Size (..))
import qualified System.Console.ANSI        as Ansi
import qualified System.IO                  as IO


--------------------------------------------------------------------------------
data Cell = Cell [Ansi.SGR] Char deriving (Cell -> Cell -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Cell -> Cell -> Bool
$c/= :: Cell -> Cell -> Bool
== :: Cell -> Cell -> Bool
$c== :: Cell -> Cell -> Bool
Eq, Int -> Cell -> ShowS
[Cell] -> ShowS
Cell -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Cell] -> ShowS
$cshowList :: [Cell] -> ShowS
show :: Cell -> String
$cshow :: Cell -> String
showsPrec :: Int -> Cell -> ShowS
$cshowsPrec :: Int -> Cell -> ShowS
Show)


--------------------------------------------------------------------------------
type Matrix = V.Vector Cell


--------------------------------------------------------------------------------
emptyCell :: Cell
emptyCell :: Cell
emptyCell = [SGR] -> Char -> Cell
Cell [] Char
' '


--------------------------------------------------------------------------------
docToMatrix :: Size -> Doc -> Matrix
docToMatrix :: Size -> Doc -> Matrix
docToMatrix Size
size Doc
doc = forall a. (forall s. ST s (MVector s a)) -> Vector a
V.create forall a b. (a -> b) -> a -> b
$ do
    MVector s Cell
matrix <- forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (MVector (PrimState m) a)
VM.replicate (Size -> Int
sRows Size
size forall a. Num a => a -> a -> a
* Size -> Int
sCols Size
size) Cell
emptyCell
    forall {f :: * -> *}.
PrimMonad f =>
MVector (PrimState f) Cell -> Int -> Int -> [Chunk] -> f ()
go MVector s Cell
matrix Int
0 Int
0 forall a b. (a -> b) -> a -> b
$ Doc -> [Chunk]
docToChunks Doc
doc
    forall (f :: * -> *) a. Applicative f => a -> f a
pure MVector s Cell
matrix
  where
    go :: MVector (PrimState f) Cell -> Int -> Int -> [Chunk] -> f ()
go MVector (PrimState f) Cell
_ Int
_ Int
_ []                                      = forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
    go MVector (PrimState f) Cell
_ Int
y Int
_ [Chunk]
_  | Int
y forall a. Ord a => a -> a -> Bool
>= Size -> Int
sRows Size
size                    = forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
    go MVector (PrimState f) Cell
r Int
y Int
_ (Chunk
NewlineChunk : [Chunk]
cs)                     = MVector (PrimState f) Cell -> Int -> Int -> [Chunk] -> f ()
go MVector (PrimState f) Cell
r (Int
y forall a. Num a => a -> a -> a
+ Int
1) Int
0 [Chunk]
cs
    go MVector (PrimState f) Cell
r Int
y Int
x [Chunk]
cs | Int
x forall a. Ord a => a -> a -> Bool
> Size -> Int
sCols Size
size                     = MVector (PrimState f) Cell -> Int -> Int -> [Chunk] -> f ()
go MVector (PrimState f) Cell
r (Int
y forall a. Num a => a -> a -> a
+ Int
1) Int
0 [Chunk]
cs
    go MVector (PrimState f) Cell
r Int
y Int
x (ControlChunk Control
ClearScreenControl  : [Chunk]
cs) = MVector (PrimState f) Cell -> Int -> Int -> [Chunk] -> f ()
go MVector (PrimState f) Cell
r Int
y Int
x [Chunk]
cs  -- ?
    go MVector (PrimState f) Cell
r Int
_ Int
x (ControlChunk (GoToLineControl Int
y) : [Chunk]
cs) = MVector (PrimState f) Cell -> Int -> Int -> [Chunk] -> f ()
go MVector (PrimState f) Cell
r Int
y Int
x [Chunk]
cs
    go MVector (PrimState f) Cell
r Int
y Int
x (StringChunk [SGR]
_      []      : [Chunk]
cs)       = MVector (PrimState f) Cell -> Int -> Int -> [Chunk] -> f ()
go MVector (PrimState f) Cell
r Int
y Int
x [Chunk]
cs
    go MVector (PrimState f) Cell
r Int
y Int
x (StringChunk [SGR]
codes (Char
z : String
zs) : [Chunk]
cs)       = do
        forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> Int -> a -> m ()
VM.write MVector (PrimState f) Cell
r (Int
y forall a. Num a => a -> a -> a
* Size -> Int
sCols Size
size forall a. Num a => a -> a -> a
+ Int
x) ([SGR] -> Char -> Cell
Cell [SGR]
codes Char
z)
        MVector (PrimState f) Cell -> Int -> Int -> [Chunk] -> f ()
go MVector (PrimState f) Cell
r Int
y (Int
x forall a. Num a => a -> a -> a
+ Char -> Int
wcwidth Char
z) ([SGR] -> String -> Chunk
StringChunk [SGR]
codes String
zs forall a. a -> [a] -> [a]
: [Chunk]
cs)


--------------------------------------------------------------------------------
hPutMatrix :: IO.Handle -> Size -> Matrix -> IO ()
hPutMatrix :: Handle -> Size -> Matrix -> IO ()
hPutMatrix Handle
h Size
size Matrix
matrix = Int -> Int -> Int -> [SGR] -> IO ()
go Int
0 Int
0 Int
0 []
  where
    go :: Int -> Int -> Int -> [SGR] -> IO ()
go !Int
y !Int
x !Int
empties [SGR]
prevCodes
        | Int
x forall a. Ord a => a -> a -> Bool
>= Size -> Int
sCols Size
size     = Handle -> String -> IO ()
IO.hPutStrLn Handle
h String
"" forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Int -> Int -> Int -> [SGR] -> IO ()
go (Int
y forall a. Num a => a -> a -> a
+ Int
1) Int
0 Int
0 [SGR]
prevCodes
        | Int
y forall a. Ord a => a -> a -> Bool
>= Size -> Int
sRows Size
size     = Handle -> [SGR] -> IO ()
Ansi.hSetSGR Handle
h [SGR
Ansi.Reset]
        -- Try to not print empty things (e.g. fill the screen with spaces) as
        -- an optimization.  Instead, store the number of empties and print them
        -- when something actually follows.
        | Cell
cell forall a. Eq a => a -> a -> Bool
== Cell
emptyCell   = do
            forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (forall (t :: * -> *) a. Foldable t => t a -> Bool
null [SGR]
prevCodes) forall a b. (a -> b) -> a -> b
$ Handle -> [SGR] -> IO ()
Ansi.hSetSGR Handle
h [SGR
Ansi.Reset]
            Int -> Int -> Int -> [SGR] -> IO ()
go Int
y (Int
x forall a. Num a => a -> a -> a
+ Int
1) (Int
empties forall a. Num a => a -> a -> a
+ Int
1) []
        | Bool
otherwise           = do
            forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (Int
empties forall a. Eq a => a -> a -> Bool
== Int
0) forall a b. (a -> b) -> a -> b
$ Handle -> String -> IO ()
IO.hPutStr Handle
h (forall a. Int -> a -> [a]
replicate Int
empties Char
' ')
            forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when ([SGR]
prevCodes forall a. Eq a => a -> a -> Bool
/= [SGR]
codes) forall a b. (a -> b) -> a -> b
$
                Handle -> [SGR] -> IO ()
Ansi.hSetSGR Handle
h (SGR
Ansi.Reset forall a. a -> [a] -> [a]
: forall a. [a] -> [a]
reverse [SGR]
codes)
            Handle -> String -> IO ()
IO.hPutStr Handle
h [Char
c]
            Int -> Int -> Int -> [SGR] -> IO ()
go Int
y (Int
x forall a. Num a => a -> a -> a
+ Char -> Int
wcwidth Char
c) Int
0 [SGR]
codes
      where
        cell :: Cell
cell@(Cell [SGR]
codes Char
c) = Matrix
matrix forall a. Vector a -> Int -> a
V.! (Int
y forall a. Num a => a -> a -> a
* Size -> Int
sCols Size
size forall a. Num a => a -> a -> a
+ Int
x)