module Wrench.Config (
    Config (..),
    readConfig,
    ensureSeed,
    executionStatsReport,
    withExecutionStats,
) where

import Data.Aeson (FromJSON (..), Value (..), genericParseJSON, rejectUnknownFields)
import Data.Aeson.Casing (aesonDrop, snakeCase)
import Data.Default
import Data.Yaml (decodeFileEither, prettyPrintParseException)
import Relude
import Relude.Extra
import Relude.Unsafe qualified as Unsafe
import System.Random qualified as Random
import Wrench.Report

throwE :: (Monad m) => e -> ExceptT e m a
throwE :: forall (m :: * -> *) e a. Monad m => e -> ExceptT e m a
throwE = m (Either e a) -> ExceptT e m a
forall e (m :: * -> *) a. m (Either e a) -> ExceptT e m a
ExceptT (m (Either e a) -> ExceptT e m a)
-> (e -> m (Either e a)) -> e -> ExceptT e m a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Either e a -> m (Either e a)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Either e a -> m (Either e a))
-> (e -> Either e a) -> e -> m (Either e a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. e -> Either e a
forall a b. a -> Either a b
Left

readConfig :: FilePath -> IO (Either String Config)
readConfig :: String -> IO (Either String Config)
readConfig String
path = ExceptT String IO Config -> IO (Either String Config)
forall e (m :: * -> *) a. ExceptT e m a -> m (Either e a)
runExceptT (ExceptT String IO Config -> IO (Either String Config))
-> ExceptT String IO Config -> IO (Either String Config)
forall a b. (a -> b) -> a -> b
$ do
    result <- IO (Either ParseException Config)
-> ExceptT String IO (Either ParseException Config)
forall a. IO a -> ExceptT String IO a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Either ParseException Config)
 -> ExceptT String IO (Either ParseException Config))
-> IO (Either ParseException Config)
-> ExceptT String IO (Either ParseException Config)
forall a b. (a -> b) -> a -> b
$ String -> IO (Either ParseException Config)
forall a. FromJSON a => String -> IO (Either ParseException a)
decodeFileEither String
path
    conf@Config{cMemoryMappedIo} <- case result of
        Left ParseException
e -> String -> ExceptT String IO Config
forall (m :: * -> *) e a. Monad m => e -> ExceptT e m a
throwE (String -> ExceptT String IO Config)
-> String -> ExceptT String IO Config
forall a b. (a -> b) -> a -> b
$ ParseException -> String
prettyPrintParseException ParseException
e
        Right Config
conf -> Config -> ExceptT String IO Config
forall a. a -> ExceptT String IO a
forall (m :: * -> *) a. Monad m => a -> m a
return Config
conf
    let conf' = (Config
conf Config -> Config -> Config
forall a. Semigroup a => a -> a -> a
<> Config
forall a. Default a => a
def){cMemoryMappedIoFlat = fmap flattenIoStream cMemoryMappedIo}
    return conf'

ensureSeed :: Config -> IO Config
ensureSeed :: Config -> IO Config
ensureSeed conf :: Config
conf@Config{cSeed :: Config -> Maybe Int
cSeed = Just Int
_} = Config -> IO Config
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return Config
conf
ensureSeed Config
conf = do
    gen <- IO StdGen
forall (m :: * -> *). MonadIO m => m StdGen
Random.initStdGen
    let (seed, _gen') = Random.uniformR (0, maxBound :: Int) gen
    return conf{cSeed = Just seed}

executionStatsReport :: ReportConf
executionStatsReport :: ReportConf
executionStatsReport =
    ReportConf
        { rcName :: Maybe Text
rcName = Text -> Maybe Text
forall a. a -> Maybe a
Just Text
"Overview"
        , rcSlice :: ReportSlice
rcSlice = ReportSlice
LastSlice
        , rcAssert :: Maybe Text
rcAssert = Maybe Text
forall a. Maybe a
Nothing
        , rcView :: Maybe Text
rcView =
            Text -> Maybe Text
forall a. a -> Maybe a
Just (Text -> Maybe Text) -> Text -> Maybe Text
forall a b. (a -> b) -> a -> b
$
                [Text] -> Text
forall t. IsText t "unlines" => [t] -> t
unlines
                    [ Text
"## Execution statistics"
                    , Text
"sim:instruction-count: {sim:instruction-count}"
                    , Text
"mem:accesses-total:    {mem:accesses-total} (reads {mem:reads-total} / writes {mem:writes-total})"
                    , Text
"layout:sections-size:  {layout:sections-size} (text {layout:text-sections-size} / data {layout:data-sections-size})"
                    , Text
"mem:instr-ranges:      {mem:instr-ranges}"
                    , Text
"mem:data-ranges:       {mem:data-ranges}"
                    , Text
"mem:io-ranges:         {mem:io-ranges}"
                    , Text
"{isa-specific}"
                    , Text
""
                    , Text
"## Memory layout"
                    , Text
"```"
                    , Text
"{mem:table}"
                    , Text
"```"
                    , Text
""
                    , Text
"## Memory accesses per instruction"
                    , Text
"```"
                    , Text
"{profile}"
                    , Text
"```"
                    ]
        }

withExecutionStats :: Config -> Config
withExecutionStats :: Config -> Config
withExecutionStats conf :: Config
conf@Config{Maybe [ReportConf]
cReports :: Maybe [ReportConf]
cReports :: Config -> Maybe [ReportConf]
cReports} =
    Config
conf{cReports = Just $ fromMaybe [] cReports <> [executionStatsReport]}

-----------------------------------------------------------

data Config = Config
    { Config -> Maybe Text
cName :: Maybe Text
    -- ^ Optional human-readable name for the configuration, purely descriptive.
    , Config -> Int
cLimit :: Int
    -- ^ The maximum number of instructions to execute.
    , Config -> Int
cMemorySize :: Int
    -- ^ The size of the memory in bytes.
    , Config -> Maybe (HashMap String [Input])
cMemoryMappedIo :: Maybe (HashMap String [Input])
    -- ^ Optional memory-mapped IO configuration, mapping stream address (decimal or hex format) to lists of inputs.
    , Config -> Maybe (IntMap ([Int], [Int]))
cMemoryMappedIoFlat :: Maybe (IntMap ([Int], [Int]))
    -- ^ (generated) Flattened memory-mapped IO configuration, mapping addresses to pairs of input and output lists.
    , Config -> Maybe [ReportConf]
cReports :: Maybe [ReportConf]
    -- ^ Optional list of report configurations.
    , Config -> Maybe Int
cSeed :: Maybe Int
    -- ^ Optional seed for random number generation.
    , Config -> Maybe Bool
cZeroMemoryInit :: Maybe Bool
    }
    deriving ((forall x. Config -> Rep Config x)
-> (forall x. Rep Config x -> Config) -> Generic Config
forall x. Rep Config x -> Config
forall x. Config -> Rep Config x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cfrom :: forall x. Config -> Rep Config x
from :: forall x. Config -> Rep Config x
$cto :: forall x. Rep Config x -> Config
to :: forall x. Rep Config x -> Config
Generic, Int -> Config -> ShowS
[Config] -> ShowS
Config -> String
(Int -> Config -> ShowS)
-> (Config -> String) -> ([Config] -> ShowS) -> Show Config
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Config -> ShowS
showsPrec :: Int -> Config -> ShowS
$cshow :: Config -> String
show :: Config -> String
$cshowList :: [Config] -> ShowS
showList :: [Config] -> ShowS
Show)

instance Default Config where
    def :: Config
def =
        Config
            { cName :: Maybe Text
cName = Maybe Text
forall a. Maybe a
Nothing
            , cLimit :: Int
cLimit = Int
1000
            , cMemorySize :: Int
cMemorySize = Int
512
            , cMemoryMappedIo :: Maybe (HashMap String [Input])
cMemoryMappedIo = Maybe (HashMap String [Input])
forall a. Maybe a
Nothing
            , cMemoryMappedIoFlat :: Maybe (IntMap ([Int], [Int]))
cMemoryMappedIoFlat = Maybe (IntMap ([Int], [Int]))
forall a. Maybe a
Nothing
            , cReports :: Maybe [ReportConf]
cReports =
                [ReportConf] -> Maybe [ReportConf]
forall a. a -> Maybe a
Just
                    [ ReportConf
                        { rcName :: Maybe Text
rcName = Text -> Maybe Text
forall a. a -> Maybe a
Just Text
"Executed Instruction Log"
                        , rcSlice :: ReportSlice
rcSlice = ReportSlice
AllSlice
                        , rcAssert :: Maybe Text
rcAssert = Maybe Text
forall a. Maybe a
Nothing
                        , rcView :: Maybe Text
rcView = Text -> Maybe Text
forall a. a -> Maybe a
Just Text
"{pc}: {instruction:next} {pc:label}\n"
                        }
                    ]
            , cSeed :: Maybe Int
cSeed = Maybe Int
forall a. Maybe a
Nothing
            , cZeroMemoryInit :: Maybe Bool
cZeroMemoryInit = Bool -> Maybe Bool
forall a. a -> Maybe a
Just Bool
False
            }

instance Semigroup Config where
    Config
a <> :: Config -> Config -> Config
<> Config
b =
        Config
            { cName :: Maybe Text
cName = Config -> Maybe Text
cName Config
a Maybe Text -> Maybe Text -> Maybe Text
forall a. Maybe a -> Maybe a -> Maybe a
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> Config -> Maybe Text
cName Config
b
            , cMemorySize :: Int
cMemorySize = Config -> Int
cMemorySize Config
a
            , cMemoryMappedIo :: Maybe (HashMap String [Input])
cMemoryMappedIo = Config -> Maybe (HashMap String [Input])
cMemoryMappedIo Config
a Maybe (HashMap String [Input])
-> Maybe (HashMap String [Input]) -> Maybe (HashMap String [Input])
forall a. Maybe a -> Maybe a -> Maybe a
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> Config -> Maybe (HashMap String [Input])
cMemoryMappedIo Config
b
            , cMemoryMappedIoFlat :: Maybe (IntMap ([Int], [Int]))
cMemoryMappedIoFlat = Config -> Maybe (IntMap ([Int], [Int]))
cMemoryMappedIoFlat Config
a Maybe (IntMap ([Int], [Int]))
-> Maybe (IntMap ([Int], [Int])) -> Maybe (IntMap ([Int], [Int]))
forall a. Maybe a -> Maybe a -> Maybe a
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> Config -> Maybe (IntMap ([Int], [Int]))
cMemoryMappedIoFlat Config
b
            , cLimit :: Int
cLimit = Config -> Int
cLimit Config
a
            , cReports :: Maybe [ReportConf]
cReports = Config -> Maybe [ReportConf]
cReports Config
a Maybe [ReportConf] -> Maybe [ReportConf] -> Maybe [ReportConf]
forall a. Maybe a -> Maybe a -> Maybe a
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> Config -> Maybe [ReportConf]
cReports Config
b
            , cSeed :: Maybe Int
cSeed = Config -> Maybe Int
cSeed Config
a Maybe Int -> Maybe Int -> Maybe Int
forall a. Maybe a -> Maybe a -> Maybe a
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> Config -> Maybe Int
cSeed Config
b
            , cZeroMemoryInit :: Maybe Bool
cZeroMemoryInit = Config -> Maybe Bool
cZeroMemoryInit Config
a Maybe Bool -> Maybe Bool -> Maybe Bool
forall a. Maybe a -> Maybe a -> Maybe a
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> Config -> Maybe Bool
cZeroMemoryInit Config
b
            }

instance FromJSON Config where
    parseJSON :: Value -> Parser Config
parseJSON = Options -> Value -> Parser Config
forall a.
(Generic a, GFromJSON Zero (Rep a)) =>
Options -> Value -> Parser a
genericParseJSON (Options -> Value -> Parser Config)
-> Options -> Value -> Parser Config
forall a b. (a -> b) -> a -> b
$ (Int -> ShowS -> Options
aesonDrop Int
1 ShowS
snakeCase){rejectUnknownFields = True}

-----------------------------------------------------------

data Input = Num Int | Chars [Int] String
    deriving (Int -> Input -> ShowS
[Input] -> ShowS
Input -> String
(Int -> Input -> ShowS)
-> (Input -> String) -> ([Input] -> ShowS) -> Show Input
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Input -> ShowS
showsPrec :: Int -> Input -> ShowS
$cshow :: Input -> String
show :: Input -> String
$cshowList :: [Input] -> ShowS
showList :: [Input] -> ShowS
Show)

instance FromJSON Input where
    parseJSON :: Value -> Parser Input
parseJSON (String Text
t) = Input -> Parser Input
forall a. a -> Parser a
forall (m :: * -> *) a. Monad m => a -> m a
return (Input -> Parser Input) -> Input -> Parser Input
forall a b. (a -> b) -> a -> b
$ [Int] -> String -> Input
Chars ((Char -> Int) -> String -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map Char -> Int
ord (String -> [Int]) -> String -> [Int]
forall a b. (a -> b) -> a -> b
$ Text -> String
forall a. ToString a => a -> String
toString Text
t) (Text -> String
forall a. ToString a => a -> String
toString Text
t)
    parseJSON (Number Scientific
n) = Input -> Parser Input
forall a. a -> Parser a
forall (m :: * -> *) a. Monad m => a -> m a
return (Input -> Parser Input) -> Input -> Parser Input
forall a b. (a -> b) -> a -> b
$ Int -> Input
Num (Scientific -> Int
forall b. Integral b => Scientific -> b
forall a b. (RealFrac a, Integral b) => a -> b
round Scientific
n) -- Int case
    parseJSON Value
_ = String -> Parser Input
forall a. String -> Parser a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"Expected a Char, String, or Int"

flattenIoStream :: HashMap String [Input] -> IntMap ([Int], [Int])
flattenIoStream :: HashMap String [Input] -> IntMap ([Int], [Int])
flattenIoStream HashMap String [Input]
memory_mapped_io =
    [Item (IntMap ([Int], [Int]))] -> IntMap ([Int], [Int])
forall l. IsList l => [Item l] -> l
fromList ([Item (IntMap ([Int], [Int]))] -> IntMap ([Int], [Int]))
-> [Item (IntMap ([Int], [Int]))] -> IntMap ([Int], [Int])
forall a b. (a -> b) -> a -> b
$ ((String, [Input]) -> Item (IntMap ([Int], [Int])))
-> [(String, [Input])] -> [Item (IntMap ([Int], [Int]))]
forall a b. (a -> b) -> [a] -> [b]
map (\(String
addr, [Input]
is) -> (String -> Int
forall a. Read a => String -> a
Unsafe.read String
addr, ([Input] -> [Int]
flatInputs [Input]
is, []))) ([(String, [Input])] -> [Item (IntMap ([Int], [Int]))])
-> [(String, [Input])] -> [Item (IntMap ([Int], [Int]))]
forall a b. (a -> b) -> a -> b
$ HashMap String [Input] -> [(String, [Input])]
forall t a b. (IsList t, Item t ~ (a, b)) => t -> [(a, b)]
toPairs HashMap String [Input]
memory_mapped_io
    where
        flatInputs :: [Input] -> [Int]
flatInputs = (Input -> [Int]) -> [Input] -> [Int]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (\case Num Int
n -> [Int
n]; Chars [Int]
ns String
_ -> [Int]
ns)