import Control.Monad ( forM_ )
import Data.List ( sortBy )
import Data.Ord ( comparing )
import Foreign
import System.Random
import Test.HUnit
import Data.Permutation
swapElems :: Storable a => Ptr a -> Int -> Int -> IO ()
swapElems ptr i j = do
x <- peekElemOff ptr i
y <- peekElemOff ptr j
pokeElemOff ptr i y
pokeElemOff ptr j x
newRandomArray :: (Random a, Storable a) => Int -> IO (Ptr a)
newRandomArray n = do
xs <- sequence $ replicate n randomIO
newArray xs
newArrayCopy :: Storable a => Int -> Ptr a -> IO (Ptr a)
newArrayCopy n x = do
y <- mallocArray n
copyArray y x n
return y
newRandomPerm :: Int -> IO (Permutation)
newRandomPerm n = do
xs <- sequence $ replicate n randomIO :: IO [Int]
let ixs = zip [0..] xs
ixs' = sortBy (comparing snd) ixs
is = (fst . unzip) ixs'
p = permutation n is
return p
testIdentity :: Test
testIdentity = TestCase $ do
let n = 30
p = identity n
x <- newRandomArray n :: IO (Ptr Double)
y <- newArrayCopy n x
applyWith (swapElems y) p
xs <- peekArray n x
ys <- peekArray n y
assertEqual ""
xs
ys
testSmallCycle :: Test
testSmallCycle = TestCase $ do
let n = 3
p = permutation 3 [ 2, 0, 1 ]
x <- newArray [ 6, 7, 8 ] :: IO (Ptr Int)
applyWith (swapElems x) p
xs <- peekArray n x
assertEqual ""
[ 8, 6, 7 ]
xs
testSmallCycleInv :: Test
testSmallCycleInv = TestCase $ do
let n = 3
p = permutation 3 [ 2, 0, 1 ]
x <- newArray [ 6, 7, 8 ] :: IO (Ptr Int)
invertWith (swapElems x) p
xs <- peekArray n x
assertEqual ""
[ 7, 8, 6 ]
xs
testSingleCycleInv :: Test
testSingleCycleInv = TestCase $ do
let n = 5
p = permutation n [ 3, 0, 1, 4, 2 ]
x <- newArray [ 0.237, 0.382, 0.818, -0.413, 0.037 ] :: IO (Ptr Double)
invertWith (swapElems x) p
xs <- peekArray n x
assertEqual ""
[ 0.382, 0.818, 0.037, 0.237, -0.413 ]
xs
testSingleCycle :: Test
testSingleCycle = TestCase $ do
let n = 5
p = permutation n [ 3, 0, 1, 4, 2 ]
x <- newArray [ 0.237, 0.382, 0.818, -0.413, 0.037 ] :: IO (Ptr Double)
applyWith (swapElems x) p
xs <- peekArray n x
assertEqual ""
[ -0.413, 0.237, 0.382, 0.037, 0.818 ]
xs
testMultiCycleInv :: Test
testMultiCycleInv = TestCase $ do
let n = 5
p = permutation n [ 2, 3, 0, 4, 1 ]
x <- newArray [ 0.237, 0.382, 0.818, -0.413, 0.037 ] :: IO (Ptr Double)
invertWith (swapElems x) p
xs <- peekArray n x
assertEqual ""
[ 0.818, 0.037, 0.237, 0.382, -0.413 ]
xs
testMultiCycle :: Test
testMultiCycle = TestCase $ do
let n = 5
p = permutation n [ 2, 3, 0, 4, 1 ]
x <- newArray [ 0.237, 0.382, 0.818, -0.413, 0.037 ] :: IO (Ptr Double)
applyWith (swapElems x) p
xs <- peekArray n x
assertEqual ""
[ 0.818, -0.413, 0.237, 0.037, 0.382 ]
xs
testApply :: Int -> Test
testApply n = TestCase $ do
x <- newRandomArray n :: IO (Ptr Double)
y <- newArrayCopy n x
p <- newRandomPerm n
applyWith (swapElems y) p
xs <- peekArray n x
ys <- peekArray n y
forM_ [0..(n-1)] $ \i ->
assertEqual ("position " ++ show i) (xs !! (apply p i)) (ys !! i)
testInvert :: Int -> Test
testInvert n = TestCase $ do
x <- newRandomArray n :: IO (Ptr Double)
y <- newArrayCopy n x
p <- newRandomPerm n
invertWith (swapElems y) p
xs <- peekArray n x
ys <- peekArray n y
forM_ [0..(n-1)] $ \i ->
assertEqual ("position " ++ show i) (xs !! i) (ys !! (apply p i))
testApplyThenInv :: Int -> Test
testApplyThenInv n = TestCase $ do
x <- newRandomArray n :: IO (Ptr Double)
y <- newArrayCopy n x
p <- newRandomPerm n
applyWith (swapElems y) p
invertWith (swapElems y) p
xs <- peekArray n x
ys <- peekArray n y
assertEqual "" xs ys
testInvThenApply :: Int -> Test
testInvThenApply n = TestCase $ do
x <- newRandomArray n :: IO (Ptr Double)
y <- newArrayCopy n x
p <- newRandomPerm n
invertWith (swapElems y) p
applyWith (swapElems y) p
xs <- peekArray n x
ys <- peekArray n y
assertEqual "" xs ys
smallTests = [ TestLabel "identity" testIdentity
, TestLabel "small cycle" testSmallCycle
, TestLabel "small cycle inverse" testSmallCycleInv
, TestLabel "single cycle" testSingleCycle
, TestLabel "single cycle inverse" testSingleCycleInv
, TestLabel "multi cycle" testMultiCycle
, TestLabel "multi cycle inverse" testMultiCycleInv
]
main =
let ns = [ 0, 1, 2, 3, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048 ]
aTests = map (\n -> TestLabel ("apply (" ++ show n ++ ")")
(testApply n)) ns
iTests = map (\n -> TestLabel ("apply (" ++ show n ++ ")")
(testInvert n)) ns
aiTests = map (\n -> TestLabel ("apply then inverse (" ++ show n ++ ")")
(testApplyThenInv n)) ns
iaTests = map (\n -> TestLabel ("inverse then apply (" ++ show n ++ ")")
(testInvThenApply n)) ns
tests = TestList $ smallTests ++ aTests ++ aiTests ++ iaTests
in runTestTT tests