{-# LANGUAGE DeriveAnyClass     #-}
{-# LANGUAGE DeriveGeneric      #-}
{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE FlexibleInstances  #-}
{-# LANGUAGE PolyKinds          #-}
{-# LANGUAGE StandaloneDeriving #-}

-----------------------------------------------------------------------------
-- |
-- Module      :  UnionFind
-- Copyright   :  (C) 2017-2019, ATS Advanced Telematic Systems GmbH
-- License     :  BSD-style (see the file LICENSE)
--
-- Maintainer  :  Stevan Andjelkovic <stevan.andjelkovic@strath.ac.uk>,
--                Momoko Hattori <momohatt10@gmail.com>
-- Stability   :  provisional
-- Portability :  non-portable (GHC extensions)
--
-- This module contains a specification for an imperative-style union/find
-- algorithm.
--
-----------------------------------------------------------------------------

module UnionFind
  ( prop_unionFindSequential
  )
  where

import           Data.Functor.Classes
                   (Eq1(..))
import           Data.IORef
                   (IORef, newIORef, readIORef, writeIORef)
import           GHC.Generics
                   (Generic, Generic1)
import           Prelude
import           Test.QuickCheck
                   (Gen, Property, arbitrary, elements, frequency,
                   shrink, (===))
import           Test.QuickCheck.Monadic
                   (monadicIO)

import           Test.StateMachine
import qualified Test.StateMachine.Types.Rank2      as Rank2
import           Test.StateMachine.Types.References
import           Test.StateMachine.Z
                   (domain, empty)

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

-- The union/find algorithm is basically an efficient way of
-- implementing an equivalence relation.
--
-- This example is borrowed from the paper
-- <http://www.cse.chalmers.se/~rjmh/Papers/QuickCheckST.ps *Testing
-- monadic code with QuickCheck*> by Koen Claessen and John Hughes
-- (2002).
--
-- Let's start with the implementation of the algorithm, which we have
-- taken from the paper.

-- Our implementation is slightly different in that we use @IORef@s
-- instead of @STRef@s (see issue #84).
data Element = Element Int (IORef Link)

data Link
  = Weight Int
  | Next Element

instance Eq Element where
  Element _ ref1 == Element _ ref2 = ref1 == ref2

instance Show Element where
  show (Element x _) = "Element " ++ show x

type Ref r = Reference (Opaque Element) r

-- We represent actions in the same way as they do in section 11 of the
-- paper.
data Command r
  = New Int
  | Find (Ref r)
  | Union (Ref r) (Ref r)
  deriving stock (Eq, Show, Generic1)
  deriving anyclass (Rank2.Functor, Rank2.Foldable, Rank2.Traversable, CommandNames)

data Response r
  = -- | New element was created.
    Created (Ref r)
    -- | Command 'Find' was successful with a return value.
  | Found (Ref r)
    -- | Command 'Union' was successful.
  | United
  deriving stock Generic1
  deriving anyclass Rank2.Foldable

deriving stock instance Show (Response Symbolic)
deriving stock instance Show (Response Concrete)

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

-- The model corresponds closely to the *relational specification* given
-- in section 12 of the paper.

newtype Model r = Model [(Ref r, Ref r)]
    deriving stock (Generic, Eq, Show)

instance ToExpr (Model Symbolic)
instance ToExpr (Model Concrete)

initModel :: Model r
initModel = Model empty

-- Find representative of 'ref'
(!) :: Eq1 r => Model r -> Ref r -> Ref r
Model m ! ref = case lookup ref m of
  Just ref' | ref == ref' -> ref
            | otherwise   -> Model m ! ref'
  Nothing                 -> error "(!): couldn't find key"

extend :: Eq1 r => Model r -> (Ref r, Ref r) -> Model r
extend (Model m) p@(ref1, ref2) =
    case lookup ref1 m of
      Just z -> -- case of 'Union' command
          let w = Model m ! ref2 in
          Model $ map (\(x, y) -> if y == z then (x, w) else (x, y)) m
      Nothing -> -- case of 'New' command
          Model $ p : m

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

precondition :: Model Symbolic -> Command Symbolic -> Logic
precondition m@(Model model) cmd = case cmd of
  New _           -> Top
  Find ref        -> ref   `member` map fst model
  Union ref1 ref2 -> (ref1 `member` map fst model) .&&
                     (ref2 `member` map fst model) .&&
                     (m ! ref1 ./= m ! ref2)

transition :: Eq1 r => Model r -> Command r -> Response r -> Model r
transition m cmd resp = case (cmd, resp) of
  (New _,           Created ref) -> m `extend` (ref, ref)
  (Find _,          Found _)     ->
      -- The equivalence relation should be the same after 'find' command.
      m
  (Union ref1 ref2, United)      ->
      -- Update the model with a new equivalence relation between ref1 and ref2.
      -- It doesn't matter whether ref1's root is pointed to ref2's root or vice versa.
      m `extend` (ref1, ref2)
  _                              -> error "transition: impossible."

postcondition :: Model Concrete -> Command Concrete -> Response Concrete -> Logic
postcondition m cmd resp = case (cmd, resp) of
  (New _,           Created ref) -> m' ! ref .== ref
  (Find _,          Found _)     -> m .== m'
  (Union ref1 ref2, United)      ->
    let z = m' ! ref1
     in (z .== m ! ref1 .|| z .== m ! ref2) .&& m' ! ref1 .== m' ! ref2
  _                              -> Bot
  where
    m' = transition m cmd resp

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

newElement :: Int -> IO Element
newElement x = do
  ref <- newIORef (Weight 1)
  return (Element x ref)

findElement :: Element -> IO Element
findElement (Element x ref) = do
  e <- readIORef ref
  case e of
    Weight _  -> return (Element x ref)
    Next next -> do
      last' <- findElement next
      writeIORef ref (Next last')
      return last'

unionElements :: Element -> Element -> IO ()
unionElements e1 e2 = do
  Element x1 ref1 <- findElement e1
  Element x2 ref2 <- findElement e2
  if ref1 == ref2
     then error "equivalent elements"
     else do
       Weight w1 <- readIORef ref1
       Weight w2 <- readIORef ref2

       if w1 <= w2
       then do
         writeIORef ref1 (Next (Element x2 ref2))
         writeIORef ref2 (Weight (w1 + w2))
       else do
         writeIORef ref2 (Next (Element x1 ref1))
         writeIORef ref1 (Weight (w1 + w2))

semantics :: Command Concrete -> IO (Response Concrete)
semantics (New x)           = Created . Reference . Concrete . Opaque <$> newElement x
semantics (Find ref)        = Found   . Reference . Concrete . Opaque <$> findElement (opaque ref)
semantics (Union ref1 ref2) = United <$ unionElements (opaque ref1) (opaque ref2)

mock :: Model Symbolic -> Command Symbolic -> GenSym (Response Symbolic)
mock _ cmd = case cmd of
  New _     -> Created <$> genSym
  Find _    -> Found <$> genSym
  Union _ _ -> pure United

generator :: Model Symbolic -> Maybe (Gen (Command Symbolic))
generator (Model m)
  | null m    = Just $ New <$> arbitrary
  | otherwise = Just $ frequency
    [ (1, New <$> arbitrary)
    , (8, Find <$> elements (domain m))
    , (8, Union <$> elements (domain m) <*> elements (domain m))
    ]

shrinker :: Model Symbolic -> Command Symbolic -> [Command Symbolic]
shrinker _ (New n) = [ New n' | n' <- shrink n ]
shrinker _ _       = []

sm :: StateMachine Model Command IO Response
sm = StateMachine initModel transition precondition postcondition
         Nothing generator shrinker semantics mock noCleanup

prop_unionFindSequential :: Property
prop_unionFindSequential =
  forAllCommands sm Nothing $ \cmds -> monadicIO $ do
    (hist, _model, res) <- runCommands sm cmds
    prettyCommands sm hist (checkCommandNames cmds (res === Ok))