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 (
    Automaton (current, states),
    FiniteOrd,
    accessSequences,
    globalCharacterizingSet,
    inputs,
    localCharacterizingSet,
    walk,
 )
import Haal.Experiment (EquivalenceOracle (..))
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 -> Either String WpMethod
mkWpMethod cfg
    | 0 <= wpmDepth cfg = Right (WpMethod cfg)
    | otherwise = Left ("Must be 0 <= wpmDepth but got " ++ show cfg)

-- | 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 (testSuffixes (localCharacterizingSet aut st)))
            | st <- Set.toAscList (states aut)
            ]
    globalSufSize = Set.size (testSuffixes (globalCharacterizingSet aut))
    transitionCover =
        Set.fromList
            [ acc ++ [a]
            | acc <- Map.elems stateCover
            , a <- alphabetList
            ]
    -- the transitions that do not already lead to a state of the state cover
    difference = transitionCover `Set.difference` Set.fromList (Map.elems stateCover)

    -- 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, testSuffixes (localCharacterizingSet aut st)) | st <- Set.toAscList $ states aut
            ]
    globalSuf = testSuffixes (globalCharacterizingSet aut)

    transitionCover =
        [ acc ++ [a]
        | acc <- Map.elems stateCover
        , a <- Set.toList alphabet
        ]
    -- the transitions that do not already lead to a state of the state cover
    difference =
        Set.fromList transitionCover
            `Set.difference` Set.fromList (Map.elems stateCover)

    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 -> Either String RandomWpMethod
mkRandomWpMethod cfg
    | 0 <= rwpMin cfg && rwpMin cfg <= rwpExpected cfg && 0 <= rwpLimit cfg = Right (RandomWpMethod cfg)
    | otherwise =
        Left
            ( "Must be 0 <= rwpMin <= rwpExpected and 0 <= rwpLimit but got "
                ++ show cfg
            )

-- | 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, testSuffixes (localCharacterizingSet aut st)) | st <- Set.toAscList $ states aut
                ]
        globalSuf = testSuffixes (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

{- | The suffixes that test words end with: the given characterizing set, or
just the empty word when that set is empty. A hypothesis with a single state
has an empty characterizing set (there is nothing to distinguish), and without
the empty word it would get no test words at all, so it would be accepted
without any testing. Ending a test word with the empty word still checks the
outputs along the rest of the word.
-}
testSuffixes :: Set.Set [i] -> Set.Set [i]
testSuffixes w
    | Set.null w = Set.singleton []
    | otherwise = w