Support initial partition for mcrl2 format

import Mdp.Mcrl2

import Mdp.Types

import StatesFile

main :: IO ()

main = hspec $ do

convertToMcrlSpec = describe "convertToMcrl" $ do

  it "generates the correct number of states" $ do

    let mdp = Mdp 4 1 mempty

    let res = convertToMcrl mdp

    let res = convertToMcrl Nothing mdp

    res ^. numStates `shouldBe` mdp ^. numStates

  it "generates a uniform initial distribution" $ do

    let mdp = Mdp 4 1 mempty

    let res = convertToMcrl mdp

    let res = convertToMcrl Nothing mdp

    res ^.. outDistribution . each . probability `shouldBe` (replicate 4 0.25)

  it "generates the correct transition if there's only one" $ do

    let mdp = Mdp 2 1 (V.singleton (Transition 0 0 1 1.0 Nothing))

    let res = convertToMcrl mdp

    let res = convertToMcrl Nothing mdp

    let trans = Mcrl2Transition 0 "0" (V.singleton (Mcrl2PropTrans 1.0 1))

    res ^? transitions . ix 0 `shouldBe` Just trans

  it "generates one out-distribution for two transitions with same label" $ do

    let mdp = Mdp 2 1 (V.fromList [Transition 0 0 0 0.5 Nothing, Transition 0 0 1 0.5 Nothing])

    let res = convertToMcrl mdp

    let res = convertToMcrl Nothing mdp

    let trans = Mcrl2Transition 0 "0" (V.fromList [Mcrl2PropTrans 0.5 0, Mcrl2PropTrans 0.5 1])

    res ^? transitions . ix 0 `shouldBe` Just trans

  it "generates two two distinct transitions for two transitions with different label" $ do

    let mdp = Mdp 2 2 (V.fromList [Transition 0 0 1 1.0 Nothing, Transition 0 1 1 1.0 Nothing])

    let res = convertToMcrl mdp

    let res = convertToMcrl Nothing mdp

    let trans1 = Mcrl2Transition 0 "0" (V.singleton (Mcrl2PropTrans 1.0 1))

    let trans2 = Mcrl2Transition 0 "1" (V.singleton (Mcrl2PropTrans 1.0 1))

    res ^.. transitions . traverse `shouldBe` [trans1, trans2]

  it "correctly models the initial partition" $ do

    let mdp = Mdp 2 2 (V.fromList [Transition 0 0 1 1.0 Nothing, Transition 0 1 1 1.0 Nothing])

    let part = Partition 2 (V.fromList [0, 1])

    let res = convertToMcrl (Just part) mdp

    let trans = [ Mcrl2Transition 0 "0" (V.singleton (Mcrl2PropTrans 1.0 1))

                , Mcrl2Transition 0 "1" (V.singleton (Mcrl2PropTrans 1.0 1))

                , Mcrl2Transition 2 "i0" (V.singleton (Mcrl2PropTrans 1.0 0))

                , Mcrl2Transition 3 "i1" (V.singleton (Mcrl2PropTrans 1.0 1))

                ]

    res ^.. transitions . traverse `shouldBe` trans

mcrl2BSpec :: Spec

mcrl2BSpec = describe "mcrl2BSpec" $ do

  it "works for an example" $ do

    let mdp = Mdp 2 2 (V.fromList [Transition 0 0 1 1.0 Nothing, Transition 0 1 1 1.0 Nothing])

    let res = "des (0 1/2 1,2,2)\n(0,\"0\",1)\n(0,\"1\",1)\n" -- FIXME Rationals

    (Build.toLazyText (mcrl2B (convertToMcrl mdp)) ^. strict) `shouldBe` res

    let res = "des (0 1/2 1,2,2)\n(0,\"0\",1)\n(0,\"1\",1)\n"

    (Build.toLazyText (mcrl2B (convertToMcrl Nothing mdp)) ^. strict) `shouldBe` res

  it "works for another example" $ do

    let mdp = Mdp 2 1 (V.fromList [Transition 0 0 0 0.5 Nothing, Transition 0 0 1 0.5 Nothing])

    let res = "des (0 1/2 1,1,2)\n(0,\"0\",0 1/2 1)\n"

    (Build.toLazyText (mcrl2B (convertToMcrl mdp)) ^. strict) `shouldBe` res

    (Build.toLazyText (mcrl2B (convertToMcrl Nothing mdp)) ^. strict) `shouldBe` res

  it "works for an example with initial partition" $ do

    let mdp = Mdp 2 2 (V.fromList [Transition 0 0 1 1.0 Nothing, Transition 0 1 1 1.0 Nothing])

    let part = Partition 2 (V.fromList [0, 1])

    let res = "des (0 1/4 1 1/4 2 1/4 3,4,4)\n(0,\"0\",1)\n(0,\"1\",1)\n(2,\"i0\",0)\n(3,\"i1\",1)\n"

    (Build.toLazyText (mcrl2B (convertToMcrl (Just part) mdp)) ^. strict) `shouldBe` res

-- FIXME Support partition

mcrl2B :: Maybe Partition -> Mdp Rational -> Build.Builder

mcrl2B _ = Mcrl2.mcrl2B . Mcrl2.convertToMcrl

mcrl2B = ((.) . (.)) Mcrl2.mcrl2B Mcrl2.convertToMcrl

import           Data.Ord

import           Control.Arrow                  ((&&&))

import           Data.Ratio

import           Data.Maybe

import           Lens.Micro.Platform

import           Data.Vector                    ( Vector )

import qualified Data.Text.Lazy.Builder.Int    as Build

import           Mdp.Types

import           StatesFile                     (Partition(..))

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

-- Builder for .aut-Format used in mcrl2

makeFields ''Mcrl2PropTrans

-- TODO Fix initial state: For now the initial state is always 0, but that

-- doesn't correspond to the actual prism model. Maybe we should provide a

-- uniform initial distribution over all states.

convertToMcrl :: Mdp Rational -> Mcrl2Mdp

convertToMcrl mdp = Mcrl2Mdp (mdp ^. numStates) (convertTransitionsToMcrl (mdp^.transitions)) initDistri

convertToMcrl :: Maybe Partition -> Mdp Rational -> Mcrl2Mdp

convertToMcrl maybePart mdp = Mcrl2Mdp stateCount ts initDistri

  where

    partitionTransitions = maybe mempty convertPartition maybePart

    ts = convertTransitionsToMcrl (mdp^.transitions) <> partitionTransitions

    stateCount = (mdp ^.numStates) * (if isJust maybePart then 2 else 1)

    initDistri = V.fromList $

      map (Mcrl2PropTrans (1.0/(fromIntegral $ mdp ^. numStates)))

          [0 .. mdp ^. numStates - 1]

      map (Mcrl2PropTrans (1%fromIntegral stateCount)) [0 .. stateCount - 1]

-- | Introduce a new state for each existing state with one edge between them,

-- labeled with the block number.

convertPartition :: Partition -> Vector Mcrl2Transition

convertPartition part = V.imap mkPartTrans (stateAssignment part)

  where

    len = length (stateAssignment part)

    mkPartTrans src blk = Mcrl2Transition (len + src) (mkPartLabel blk) (mkPartDist src)

    mkPartLabel blk = "i" <> T.pack (show blk)

    mkPartDist src = V.singleton (Mcrl2PropTrans 1 src)

convertTransitionsToMcrl :: Vector (Transition Rational) -> Vector Mcrl2Transition

convertTransitionsToMcrl trans = mkMcrl2Tans <$> partitionVector (view source &&& view choice) sorted