Matrix.hs

{-# LANGUAGE DataKinds #-}

module Tests.Extra.Semigroup.Matrix (tests) where

import AtCoder.Extra.Semigroup.Matrix qualified as Mat
import AtCoder.ModInt qualified as M
import Data.Semigroup (stimes)
import Data.Vector.Unboxed qualified as VU
import GHC.TypeNats (KnownNat)
import Test.QuickCheck.Classes qualified as QCC
import Test.Tasty
import Test.Tasty.QuickCheck qualified as QC
import Tests.Util (laws)

-- TODO: (const True) should be removed

-- orphan instance
instance (QC.Arbitrary a, VU.Unbox a) => QC.Arbitrary (Mat.Matrix a) where
  -- for simplicity, make a 33x33 matrix
  arbitrary = do
    let n = 33
    vec <- VU.fromList <$> QC.vectorOf (n * n) (QC.arbitrary @a)
    pure $ Mat.Matrix n n vec

-- orphan instance
instance (KnownNat p) => QC.Arbitrary (M.ModInt p) where
  arbitrary = M.new <$> QC.arbitrary

prop_mulToCol :: QC.Gen QC.Property
prop_mulToCol = do
  h <- QC.chooseInt (1, 16)
  w <- QC.chooseInt (1, 16)
  vec <- VU.fromList <$> QC.vectorOf (h * w) (QC.arbitrary @Int)
  let mat = Mat.new h w vec
  col <- VU.fromList <$> QC.vectorOf w (QC.arbitrary @Int)
  let lhs = Mat.mulToCol mat col
  let rhs = Mat.vecM $ Mat.mul mat (Mat.new w 1 col)
  pure  $ lhs QC.=== rhs

m :: Int
m = 998244353

prop_pow :: QC.Small Int -> QC.NonNegative Int -> QC.Gen QC.Property
prop_pow (QC.Small n) (QC.NonNegative k) = do
  vec <- VU.fromList <$> QC.vectorOf (n * n) (QC.chooseInt (0, m - 1))
  let mat = Mat.new n n vec
  if k == 0
    then pure $ Mat.pow k mat QC.=== Mat.ident n
    else pure $ Mat.pow k mat QC.=== stimes k mat

tests :: [TestTree]
tests =
  [ QC.testProperty "mulToCol" prop_mulToCol,
    laws @(Mat.Matrix (M.ModInt 998244353))
      [ QCC.semigroupLaws
      ],
    QC.testProperty "pow" prop_pow
  ]