BlackBox.hs

{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE ScopedTypeVariables #-}

{- | This module defines the BlackBox type class as well as the Automaton and SUL
sub classes.
-}
module Haal.BlackBox (
    Automaton (..),
    SUL (..),
    StateID,
    Finite,
    FiniteEq,
    FiniteOrd,
    inputs,
    outputs,
    walk,
    stepPure,
    walkPure,
    resetPure,
    initial,
    distinguish,
    accessSequences,
    localCharacterizingSet,
    globalCharacterizingSet,
    reachable,
)
where

import Control.Monad.Identity (Identity, runIdentity)
import qualified Data.Bifunctor as Bif
import qualified Data.List as List
import qualified Data.Map as Map
import qualified Data.Set as Set

{- | The 'StateID' type is an alias for an integer that represents the state of the automaton.
 - It is used as a default type for the state of learned automata.
-}
type StateID = Int

{- | The 'SUL' type class defines the basic interface for a black box automaton.
It provides methods to step through the automaton and retrieve the current state.
It also requires a monad m, that may be 'Identity' in case of a pure SUL, or 'IO'
in case of an external program that performs IO.
-}
class (Monad m) => SUL sul m where
    step :: sul i o -> i -> m (sul i o, o)
    reset :: sul i o -> m (sul i o)

-- | Finite is an alias for (Enum, Bounded).
type Finite i = (Enum i, Bounded i)
-- | FiniteEq is an alias for (Eq, Finite).
type FiniteEq i = (Eq i, Finite i)
-- | FiniteOrd is an alias for (Ord, Bounded).
type FiniteOrd i = (Ord i, Finite i)

-- | Generalization of 'step' that operates on a list of inputs.
walk :: (SUL sul m) => sul i o -> [i] -> m (sul i o, [o])
walk sul [] = pure (sul, [])
walk sul (x : xs) = do
    (sul', o) <- step sul x
    (sul'', os) <- walk sul' xs
    pure (sul'', o : os)

-- | Return a Set containing only the valid inputs of the SUL.
inputs :: (FiniteOrd i) => sul i o -> Set.Set i
inputs _ = Set.fromList [minBound .. maxBound]

-- | Return a Set containing only the valid outputs of the SUL.
outputs :: (FiniteOrd o) => sul i o -> Set.Set o
outputs _ = Set.fromList [minBound .. maxBound]

{- | The 'Automaton' type class extends the 'SUL' type class and adds
support for automata operations. Automatons are models, not programs,
so they are pure and operate in the Identity monad.
-}
class (SUL (aut s) Identity) => Automaton aut s where
    transitions ::
        (FiniteOrd i, FiniteOrd s) =>
        aut s i o ->
        Map.Map (s, i) (s, o)
    states :: (FiniteOrd s) => aut s i o -> Set.Set s
    current :: aut s i o -> s
    update :: aut s i o -> s -> aut s i o

-- | Pure instance of 'step'.
stepPure :: (SUL sul Identity) => sul i o -> i -> (sul i o, o)
stepPure sul i = runIdentity (step sul i)

-- | Pure instance of 'walk'.
walkPure :: (SUL sul Identity) => sul i o -> [i] -> (sul i o, [o])
walkPure sul i = runIdentity (walk sul i)

-- | Pure instance of 'reset'.
resetPure :: (SUL sul Identity) => sul i o -> sul i o
resetPure sul = runIdentity (reset sul)

-- | Return the initial state of an automaton.
initial :: (Automaton aut s) => aut s i o -> s
initial = current . resetPure

-- | Return the set of reachable states of an automaton.
reachable :: forall s i o aut. (Automaton aut s, Ord s, FiniteOrd i) => aut s i o -> Set.Set s
reachable aut = bfs [initial aut] $ Set.singleton (initial aut)
  where
    alphabet = inputs aut

    bfs :: [s] -> Set.Set s -> Set.Set s
    bfs [] visited = visited
    bfs (st : queue) visited = bfs queue' visited'
      where
        aut' = update aut st
        neighbours = Set.map (current . fst . stepPure aut') alphabet
        visited' = visited `Set.union` neighbours
        queue' = Set.toList (neighbours `Set.difference` visited) ++ queue

-- | Returns a map containing the shortest sequence to access each reachable state from the initial state.
accessSequences ::
    forall s i o aut.
    (Automaton aut s, FiniteOrd i, Ord s) =>
    aut s i o ->
    Map.Map s [i]
accessSequences aut = bfs [(initialSt, [])] (Set.singleton initialSt) (Map.singleton initialSt [])
  where
    alphabet = Set.toList (inputs aut)
    initialSt = initial aut

    bfs :: [(s, [i])] -> Set.Set s -> Map.Map s [i] -> Map.Map s [i]
    bfs [] _ acc = Map.map List.reverse acc
    bfs ((_, prefix) : rest) visited acc =
        bfs (rest ++ newQueue) newVisited newMap
      where
        mo = fst $ walkPure (resetPure aut) (reverse prefix)
        successors =
            [ (nextState, input : prefix)
            | input <- alphabet
            , let nextState = current . fst $ stepPure mo input
            , nextState `Set.notMember` visited
            ]

        newMap = foldr (uncurry Map.insert) acc successors
        newVisited = foldr (Set.insert . fst) visited successors
        newQueue = successors

{- | Returns an input sequence that distinguishes the given states in
the given automaton.
-}
distinguish ::
    ( Automaton aut s
    , FiniteOrd i
    , Ord s
    , Eq o
    ) =>
    aut s i o ->
    s ->
    s ->
    [i]
distinguish m s1 s2 = explore Map.empty [(s1, s2, [])]
  where
    alphabet = Set.toList (inputs m)

    explore _ [] = []
    explore visited ((q1, q2, prefix) : queue)
        | Just symbol <- discrepancy = reverse (symbol : prefix)
        | otherwise = explore newVisited (queue ++ newQueue)
      where
        newVisited = Map.insert (q1, q2) prefix visited
        mo1 = update m q1
        mo2 = update m q2

        (nextStates1, outputs1) = unzip $ map (stepAndCurrent mo1) alphabet
        (nextStates2, outputs2) = unzip $ map (stepAndCurrent mo2) alphabet

        discrepancy = snd <$> List.find fst (zip (zipWith (/=) outputs1 outputs2) alphabet)

        appended = map (: prefix) alphabet

        toBeVisited = Map.fromList $ zip (zip nextStates1 nextStates2) appended

        newQueue = [(s1', s2', p) | ((s1', s2'), p) <- Map.toList toBeVisited, (s1', s2') `Map.notMember` visited]

    stepAndCurrent mo i = Bif.first current (stepPure mo i)

{- | Returns a set of lists of inputs that can be used to distinguish between the given state and
 - any other state of the automaton.
-}
localCharacterizingSet ::
    ( Automaton aut s
    , FiniteOrd i
    , FiniteOrd s
    , Eq o
    ) =>
    aut s i o ->
    s ->
    Set.Set [i]
localCharacterizingSet m s = Set.fromList [d s sx | sx <- Set.toList $ states m, s /= sx]
  where
    d = distinguish m

{- | Returns a set of lists of inputs that can be used to distinguish between any two different states
of the automaton.
-}
globalCharacterizingSet ::
    ( Automaton aut s
    , FiniteOrd i
    , FiniteOrd s
    , Eq o
    ) =>
    aut s i o ->
    Set.Set [i]
globalCharacterizingSet m = Set.fromList [d s1 s2 | s1 <- sts, s2 <- sts, s1 < s2]
  where
    sts = Set.toList $ states m
    d = distinguish m