--------------------------------------------------------------------------------
{-# 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
(Cell -> Cell -> Bool) -> (Cell -> Cell -> Bool) -> Eq Cell
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Cell -> Cell -> Bool
== :: Cell -> Cell -> Bool
$c/= :: Cell -> Cell -> Bool
/= :: Cell -> Cell -> Bool
Eq, Int -> Cell -> ShowS
[Cell] -> ShowS
Cell -> String
(Int -> Cell -> ShowS)
-> (Cell -> String) -> ([Cell] -> ShowS) -> Show Cell
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Cell -> ShowS
showsPrec :: Int -> Cell -> ShowS
$cshow :: Cell -> String
show :: Cell -> String
$cshowList :: [Cell] -> ShowS
showList :: [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 Int
rows Int
cols) Doc
doc = (forall s. ST s (MVector s Cell)) -> Matrix
forall a. (forall s. ST s (MVector s a)) -> Vector a
V.create ((forall s. ST s (MVector s Cell)) -> Matrix)
-> (forall s. ST s (MVector s Cell)) -> Matrix
forall a b. (a -> b) -> a -> b
$ do
    MVector s Cell
matrix <- Int -> Cell -> ST s (MVector (PrimState (ST s)) Cell)
forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (MVector (PrimState m) a)
VM.replicate (Int
rows Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
cols) Cell
emptyCell
    MVector (PrimState (ST s)) Cell -> Int -> Int -> [Chunk] -> ST s ()
forall {f :: * -> *}.
PrimMonad f =>
MVector (PrimState f) Cell -> Int -> Int -> [Chunk] -> f ()
go MVector s Cell
MVector (PrimState (ST s)) Cell
matrix Int
0 Int
0 ([Chunk] -> ST s ()) -> [Chunk] -> ST s ()
forall a b. (a -> b) -> a -> b
$ Doc -> [Chunk]
docToChunks Doc
doc
    MVector s Cell -> ST s (MVector s Cell)
forall a. a -> ST s a
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
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
_ Int
_ Int
_ []                                      = () -> f ()
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
    go MVector (PrimState f) Cell
_ Int
y Int
_ [Chunk]
_  | Int
y Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
rows                          = () -> f ()
forall a. a -> f a
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 Int -> Int -> Int
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 chunks :: [Chunk]
chunks@(StringChunk [SGR]
codes (Char
z : String
zs) : [Chunk]
cs)
        | Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
w Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
cols = MVector (PrimState f) Cell -> Int -> Int -> [Chunk] -> f ()
go MVector (PrimState f) Cell
r (Int
y Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int
0 [Chunk]
chunks
        | Bool
otherwise    = do
            MVector (PrimState f) Cell -> Int -> Cell -> f ()
forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> Int -> a -> m ()
VM.write MVector (PrimState f) Cell
r (Int
y Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
cols Int -> Int -> Int
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 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Char -> Int
wcwidth Char
z) ([SGR] -> String -> Chunk
StringChunk [SGR]
codes String
zs Chunk -> [Chunk] -> [Chunk]
forall a. a -> [a] -> [a]
: [Chunk]
cs)
      where
        w :: Int
w = Char -> Int
wcwidth Char
z


--------------------------------------------------------------------------------
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 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Size -> Int
sCols Size
size     = Handle -> String -> IO ()
IO.hPutStrLn Handle
h String
"" IO () -> IO () -> IO ()
forall a b. IO a -> IO b -> IO b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Int -> Int -> Int -> [SGR] -> IO ()
go (Int
y Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int
0 Int
0 [SGR]
prevCodes
        | Int
y Int -> Int -> Bool
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 Cell -> Cell -> Bool
forall a. Eq a => a -> a -> Bool
== Cell
emptyCell   = do
            Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless ([SGR] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [SGR]
prevCodes) (IO () -> IO ()) -> IO () -> IO ()
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 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int
empties Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) []
        | Bool
otherwise           = do
            Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (Int
empties Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ Handle -> String -> IO ()
IO.hPutStr Handle
h (Int -> Char -> String
forall a. Int -> a -> [a]
replicate Int
empties Char
' ')
            Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when ([SGR]
prevCodes [SGR] -> [SGR] -> Bool
forall a. Eq a => a -> a -> Bool
/= [SGR]
codes) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$
                Handle -> [SGR] -> IO ()
Ansi.hSetSGR Handle
h (SGR
Ansi.Reset SGR -> [SGR] -> [SGR]
forall a. a -> [a] -> [a]
: [SGR] -> [SGR]
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 Int -> Int -> Int
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 Matrix -> Int -> Cell
forall a. Vector a -> Int -> a
V.! (Int
y Int -> Int -> Int
forall a. Num a => a -> a -> a
* Size -> Int
sCols Size
size Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
x)