TestContext.hs

{-# LANGUAGE OverloadedStrings #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}

-- |
-- SPDX-License-Identifier: BSD-3-Clause
--
-- Swarm unit tests for contexts
module TestContext where

import Control.Monad (replicateM)
import Data.Hashable
import Data.List (nub)
import Data.Map qualified as M
import Data.Yaml (decodeEither', encode)
import Swarm.Language.Context
import Swarm.Language.Var (Var)
import Swarm.Util (showT)
import Test.QuickCheck.Instances.Text ()
import Test.Tasty (TestTree, testGroup)
import Test.Tasty.HUnit (Assertion, assertBool, assertEqual, testCase)
import Test.Tasty.QuickCheck (Arbitrary (..), Gen, generate, testProperty, withMaxSuccess)

testContext :: TestTree
testContext =
  testGroup
    "Contexts"
    [ testGroup
        "Context equality"
        [ testCase "idempotence 1" $ ctxsEqual ctx1 (ctx1 <> ctx1)
        , testCase "idempotence 2" $ ctxsEqual ctx2 (ctx2 <> ctx2)
        , testCase "deletion" $ ctxsEqual ctx1 (delete "z" ctx2)
        , testCase "empty/delete" $ ctxsEqual empty (delete "x" ctx1)
        , testCase "fromMap" $ ctxsEqual ctx2 (fromMap (M.fromList [("x", 3), ("z", 6)]))
        , testCase "right bias" $ ctxsEqual ctx4 (ctx2 <> ctx3)
        , testCase "commutativity" $ ctxsEqual (ctx1 <> ctx5) (ctx5 <> ctx1)
        ]
    , testGroup
        "de/rehydrate round-trip"
        (map (\(name, ctx) -> testCase name $ hydrationRoundTrip ctx) testCtxs)
    , testGroup
        "serialize round-trip"
        (map (\(name, ctx) -> testCase name $ serializeRoundTrip ctx) testCtxs)
    , testProperty
        "no paired hash collisions"
        (withMaxSuccess 10000 (hashConsistent @Var @Int))
    , testCase
        "no hash collisions in a large pool"
        $ do
          ctxs <- generate (replicateM 100000 (arbitrary :: Gen (Ctx Var Int)))
          let m = M.fromListWith (++) (map (\ctx -> (ctxHash ctx, [unCtx ctx])) ctxs)
          assertBool "foo" $ all ((== 1) . length . nub) m
    ]
 where
  ctx1 = singleton "x" 3
  ctx2 = singleton "x" 3 <> singleton "z" 6
  ctx3 = singleton "x" 5 <> singleton "y" 7
  ctx4 = singleton "x" 5 <> singleton "y" 7 <> singleton "z" 6
  ctx5 = singleton "y" 10
  bigCtx = fromMap . M.fromList $ zip (map (("x" <>) . showT) [1 :: Int ..]) [1 .. 10000]

  testCtxs = [("empty", empty), ("ctx1", ctx1), ("ctx2", ctx2), ("ctx3", ctx3), ("ctx4", ctx4), ("ctx5", ctx5), ("large", bigCtx), ("delete", delete "y" ctx4)]

instance (Ord v, Hashable v, Arbitrary v, Hashable a, Arbitrary a) => Arbitrary (Ctx v a) where
  arbitrary = fromMap <$> arbitrary

hashConsistent :: (Eq v, Eq a) => Ctx v a -> Ctx v a -> Bool
hashConsistent ctx1 ctx2 = (ctx1 == ctx2) == (ctx1 `ctxStructEqual` ctx2)

ctxsEqual :: Ctx Var Int -> Ctx Var Int -> Assertion
ctxsEqual ctx1 ctx2 = do
  -- Contexts are compared by hash for equality
  assertEqual "hash equality" ctx1 ctx2

  -- Make sure they are also structurally equal
  assertBool "structural equality" (ctxStructEqual ctx1 ctx2)

ctxStructEqual :: (Eq a, Eq v) => Ctx v a -> Ctx v a -> Bool
ctxStructEqual (Ctx m1 _) (Ctx m2 _) = m1 == m2

hydrationRoundTrip :: Ctx Var Int -> Assertion
hydrationRoundTrip ctx = do
  case getCtx (ctxHash ctx) (rehydrate (dehydrate (toCtxMap ctx))) of
    Nothing -> fail "Failed to reconstitute dehydrated context"
    Just ctx' -> ctxsEqual ctx ctx'

serializeRoundTrip :: Ctx Var Int -> Assertion
serializeRoundTrip ctx = do
  case decodeEither' (encode (dehydrate (toCtxMap ctx))) of
    Left e -> fail $ "Failed to decode JSON-encoded context: " ++ show e
    Right c -> case getCtx (ctxHash ctx) (rehydrate c) of
      Nothing -> fail "Failed to reconstitute dehydrated context"
      Just ctx' -> ctxsEqual ctx ctx'