{-# LANGUAGE DataKinds #-}
{-# LANGUAGE TypeFamilies #-}
{-# OPTIONS_GHC -Wno-orphans #-}

-- | Edge-based bitmap trace store.
module Test.Mutagen.Tracer.Store.Bitmap
  ( -- * Bitmap trace store
    TraceStore
  , newTraceStore
  , resetTraceStore
  , saveTrace
  , printTraceStore
  )
where

import Control.Monad (forM_, replicateM_)
import Data.Array.IO (IOUArray, getBounds, newArray, readArray, writeArray)
import Data.Foldable (foldrM)
import Test.Mutagen.Tracer.Store.API (TraceStoreImpl (..))
import Test.Mutagen.Tracer.Store.Types (TraceBackend (Bitmap))
import Test.Mutagen.Tracer.Trace (Trace (..))

{-------------------------------------------------------------------------------
-- * Bitmap trace store
-------------------------------------------------------------------------------}

instance TraceStoreImpl Bitmap where
  data TraceStore Bitmap = BitmapStore Int (IOUArray Int Bool)

  type SaveTraceResult Bitmap = Int

  newTraceStore :: Int -> IO (TraceStore 'Bitmap)
newTraceStore Int
n = do
    Int -> IOUArray Int Bool -> TraceStore 'Bitmap
BitmapStore (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (IOUArray Int Bool -> TraceStore 'Bitmap)
-> IO (IOUArray Int Bool) -> IO (TraceStore 'Bitmap)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Int, Int) -> Bool -> IO (IOUArray Int Bool)
forall i. Ix i => (i, i) -> Bool -> IO (IOUArray i Bool)
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
(i, i) -> e -> m (a i e)
newArray (Int
0, (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int -> Int -> Int
forall a. Num a => a -> a -> a
* (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)) Bool
False

  resetTraceStore :: TraceStore 'Bitmap -> IO ()
resetTraceStore (BitmapStore Int
_ IOUArray Int Bool
arr) = do
    (l, u) <- IOUArray Int Bool -> IO (Int, Int)
forall i. Ix i => IOUArray i Bool -> IO (i, i)
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> m (i, i)
getBounds IOUArray Int Bool
arr
    forM_ [l .. u] $ \Int
i -> do
      IOUArray Int Bool -> Int -> Bool -> IO ()
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> e -> m ()
writeArray IOUArray Int Bool
arr Int
i Bool
False

  saveTrace :: Trace -> TraceStore 'Bitmap -> IO (SaveTraceResult 'Bitmap)
saveTrace (Trace [Int]
entries) (BitmapStore Int
n IOUArray Int Bool
arr) = do
    let edges :: [(Int, Int)]
edges = [Int] -> [Int] -> [(Int, Int)]
forall a b. [a] -> [b] -> [(a, b)]
zip (Int
0 Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: [Int]
entries) [Int]
entries
    let flipAndCount :: (Int, Int) -> Int -> IO Int
flipAndCount (Int
i, Int
j) Int
acc = do
          let idx :: Int
idx = Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
j
          b <- IOUArray Int Bool -> Int -> IO Bool
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray IOUArray Int Bool
arr Int
idx
          if b
            then return acc
            else writeArray arr idx True >> return (acc + 1)
    ((Int, Int) -> Int -> IO Int) -> Int -> [(Int, Int)] -> IO Int
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> b -> m b) -> b -> t a -> m b
foldrM (Int, Int) -> Int -> IO Int
flipAndCount Int
0 [(Int, Int)]
edges

  printTraceStore :: TraceStore 'Bitmap -> IO ()
printTraceStore (BitmapStore Int
n IOUArray Int Bool
arr) = do
    String -> IO ()
putStr 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 -> IO () -> IO ()
forall (m :: * -> *) a. Applicative m => Int -> m a -> m ()
replicateM_ Int
n (String -> IO ()
putStr 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
>> String -> IO ()
putStrLn String
"+"
    [Int] -> (Int -> IO ()) -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [Int
0 .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1] ((Int -> IO ()) -> IO ()) -> (Int -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Int
i -> do
      String -> IO ()
putStr String
"|"
      [Int] -> (Int -> IO ()) -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [Int
0 .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1] ((Int -> IO ()) -> IO ()) -> (Int -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Int
j -> do
        b <- IOUArray Int Bool -> Int -> IO Bool
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray IOUArray Int Bool
arr (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
j)
        if b then putStr "*" else putStr " "
      String -> IO ()
putStrLn String
"|"
    String -> IO ()
putStr 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 -> IO () -> IO ()
forall (m :: * -> *) a. Applicative m => Int -> m a -> m ()
replicateM_ Int
n (String -> IO ()
putStr 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
>> String -> IO ()
putStrLn String
"+"