WpMethod.hs

{-# LANGUAGE ScopedTypeVariables #-}

-- | This module implements the WpMethod.
module Haal.EquivalenceOracle.WpMethod (
    WpMethod,
    WpMethodConfig (..),
    RandomWpMethod,
    RandomWpMethodConfig (..),
    wpmethodSuiteSize,
    mkWpMethod,
    mkRandomWpMethod,
) where

import Control.Monad (replicateM)
import Control.Monad.Identity (Identity (runIdentity))
import Control.Monad.State (MonadState (state), State, runState)
import qualified Data.Map as Map
import qualified Data.Set as Set
import Haal.BlackBox
import Haal.Experiment
import System.Random (Random (randomR), StdGen)

-- | The 'WpMethodConfig' type is used to configure the Wp-method equivalence oracle.
data WpMethodConfig = WpMethodConfig
    { wpmDepth :: Int
    -- ^ The number of extra states beyond the hypothesis to account for.
    }
    deriving (Show, Eq)

-- | The 'WpMethod' type represents the Wp-method equivalence oracle.
newtype WpMethod = WpMethod WpMethodConfig deriving (Eq, Show)

-- | Constructor for a 'WpMethod' value.
mkWpMethod :: WpMethodConfig -> WpMethod
mkWpMethod = WpMethod

-- | The 'wpmethodSuiteSize' returns the number of test cases in the test suite of WpMethod.
wpmethodSuiteSize ::
    ( Automaton aut s
    , FiniteOrd i
    , FiniteOrd s
    , Eq o
    ) =>
    WpMethod ->
    aut s i o ->
    Int
wpmethodSuiteSize (WpMethod (WpMethodConfig d)) aut = firstPhaseSize + secondPhaseSize
  where
    alphabetList = Set.toList (inputs aut)
    alphabetSize = length alphabetList
    stateCover = accessSequences aut
    localSufSizes =
        Map.fromAscList
            [ (st, Set.size (localCharacterizingSet aut st))
            | st <- Set.toAscList (states aut)
            ]
    globalSufSize = Set.size (globalCharacterizingSet aut)
    transitionCover =
        Set.fromList
            [ acc ++ [a]
            | acc <- Map.elems stateCover
            , a <- alphabetList
            ]
    difference = Set.fromList (Map.elems stateCover) `Set.difference` transitionCover

    -- Closed form: |S| * |W| * (1 + |Σ| + ... + |Σ|^d)
    firstPhaseSize =
        Map.size stateCover
            * globalSufSize
            * sum [alphabetSize ^ k | k <- [0 .. d]]

    -- Walk each acc once, then enumerate middles from that state; no test strings built.
    secondPhaseSize =
        sum
            [ localSufSizes Map.! current (fst (runIdentity (walk afterAcc middle)))
            | acc <- Set.toList difference
            , let afterAcc = fst (runIdentity (walk aut acc))
            , fixed <- [0 .. d]
            , middle <- replicateM fixed alphabetList
            ]

-- | Returns the test suite for the WpMethod.
wpmethodSuite ::
    forall aut i o s.
    ( Automaton aut s
    , FiniteOrd i
    , FiniteOrd s
    , Eq o
    ) =>
    WpMethod ->
    aut s i o ->
    (WpMethod, [[i]])
wpmethodSuite wpm@(WpMethod (WpMethodConfig d)) aut = (wpm, suite)
  where
    alphabet = inputs aut
    stateCover = accessSequences aut
    localSuf =
        Map.fromAscList
            [ (st, localCharacterizingSet aut st) | st <- Set.toAscList $ states aut
            ]
    globalSuf = globalCharacterizingSet aut

    transitionCover =
        [ acc ++ [a]
        | acc <- Map.elems stateCover
        , a <- Set.toList alphabet
        ]
    difference =
        Set.fromList (Map.elems stateCover)
            `Set.difference` Set.fromList transitionCover

    firstPhase =
        concat
            [ [ acc ++ middle ++ suf
              | acc <- Map.elems stateCover
              , suf <- Set.toList globalSuf
              ]
            | fixed <- [0 .. d]
            , middle <- replicateM fixed $ Set.toList alphabet
            ]

    secondPhase =
        concat
            [ [ acc ++ middle ++ suf
              | acc <- Set.toList difference
              , suf <- Set.toList $ localSuf Map.! current (fst (runIdentity (walk aut (acc ++ middle))))
              ]
            | fixed <- [0 .. d]
            , middle <- replicateM fixed $ Set.toList alphabet
            ]

    suite = firstPhase ++ secondPhase

{- | The 'RandomWpMethodConfig' is a record data type that represents the configuration for an instance
of the Random WpMethod algorithm.
-}
data RandomWpMethodConfig = RandomWpMethodConfig
    { rwpGen :: StdGen
    -- ^ Random generator.
    , rwpExpected :: Int
    -- ^ Expected depth of random walk.
    , rwpMin :: Int
    -- ^ Minimum depth of random walk.
    , rwpLimit :: Int
    -- ^ Maximum number of queries.
    }
    deriving (Show, Eq)

-- | The 'RandomWpMethod' type is just a wrapper around the config.
newtype RandomWpMethod = RandomWpMethod RandomWpMethodConfig deriving (Show, Eq)

-- | Constructor for a 'RandomWpMethod' value.
mkRandomWpMethod :: RandomWpMethodConfig -> RandomWpMethod
mkRandomWpMethod = RandomWpMethod

-- | Return the 'RandomWpMethod' test suite.
randomWpMethodSuite ::
    forall aut i o s.
    ( Automaton aut s
    , FiniteOrd i
    , FiniteOrd s
    , Eq o
    ) =>
    RandomWpMethod ->
    aut s i o ->
    (RandomWpMethod, [[i]])
randomWpMethodSuite
    ( RandomWpMethod
            conf@RandomWpMethodConfig
                { rwpGen = g
                , rwpExpected = e
                , rwpMin = mi
                , rwpLimit = lim
                }
        )
    aut = (RandomWpMethod (conf{rwpGen = genfinal}), suite)
      where
        alphabet = inputs aut
        prefixes = accessSequences aut
        localSuf =
            Map.fromAscList
                [ (st, localCharacterizingSet aut st) | st <- Set.toAscList $ states aut
                ]
        globalSuf = globalCharacterizingSet aut

        (suite, genfinal) = runState (replicateM lim genTestCase) g

        genTestCase :: State StdGen [i]
        genTestCase = do
            prefixIdx <- state $ randomR (0, Map.size prefixes - 1)
            let (_, prefix) = prefixIdx `Map.elemAt` prefixes
            middle <- genExpectedLength
            local <- state $ randomR (False, True)
            if local
                then do
                    let curr = current $ fst (runIdentity (walk aut (prefix ++ middle)))
                        suffixSet = localSuf Map.! curr
                    suffixIdx <- state $ randomR (0, Set.size suffixSet - 1)
                    let suffix = suffixIdx `Set.elemAt` suffixSet
                    return $ prefix ++ middle ++ suffix
                else do
                    globalIdx <- state $ randomR (0, Set.size globalSuf - 1)
                    let suffix = globalIdx `Set.elemAt` globalSuf
                    return $ prefix ++ middle ++ suffix

        genExpectedLength :: State StdGen [i]
        genExpectedLength = state $ go [] mi
          where
            go :: [i] -> Int -> StdGen -> ([i], StdGen)
            go acc minim gen =
                let (continue, gen') = randomR (0.0 :: Double, 1.0) gen
                 in if minim > 0 || continue > 1 / (fromIntegral e + 1)
                        then
                            let (idx, gen'') = randomR (0, Set.size alphabet - 1) gen'
                                nextChar = idx `Set.elemAt` alphabet
                             in go (nextChar : acc) (minim - 1) gen''
                        else (acc, gen)

instance EquivalenceOracle WpMethod where
    testSuite = wpmethodSuite

instance EquivalenceOracle RandomWpMethod where
    testSuite = randomWpMethodSuite