packages feed

astar-monad-0.3.0.0: test/Spec.hs

{-# LANGUAGE QuasiQuotes #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE TypeApplications #-}
import Control.Monad.AStar
import Test.Hspec hiding (Arg, context)
import Data.Foldable
import Control.Lens hiding (Context)
import Control.Monad.State
import Control.Applicative
import Data.Monoid

data Move = U | D | L | R
    deriving (Show, Eq)

data Context =
    Context { _current   :: (Int, Int)
            , _goal      :: (Int, Int)
            , _moves     :: [Move]
            , _locations :: [(Int, Int)]
            }
    deriving (Show, Eq)

makeLenses ''Context

context :: (Int, Int) -> (Int, Int) -> Context
context start goal = Context start goal [] []

main :: IO ()
main = hspec $ do
    describe "a-star" $ do
        it "should find a solution" $ do
            (fst <$> runAStar findPoint (context (3, 6) (5, 5)))
              `shouldBe` Just (5, 5)
        it "should take the shortest path" $ do
            (view moves . snd <$> runAStar findPoint (context (4, 6) (5, 5)))
              `shouldBe` Just [R, U]
        it "should take the shortest path in long situations" $ do
            (length . view moves . snd <$> runAStar findPoint (context (4, 6) (20, 20)))
              `shouldBe` Just 30
        -- it "should properly rewind state" $ do
        --       do flip runAStar [] $ do
        --             asum [ estimate (1 :: Int) >> modify (++ [1]) >> estimate 10 >> modify (++ [10])
        --                  , estimate (2 :: Int) >> modify (++ [2]) >> done ()
        --                  ]
        it "should resolve with Nothing if branches return after updating cost" $ do
              do flip runAStar () $ (estimate @(Sum Int) 1 >> return ()) <|> return ()
            `shouldBe`
              (Nothing :: Maybe ((), ()))
        it "should resolve with solution if some branches simply return" $ do
              do flip runAStar () $ (return () <|> (estimate @(Sum Int) 1 >> done ()))
            `shouldBe`
              Just ((), ())
        it "should resolve without solution if all branches simply return" $ do
              do flip runAStar () $ (return () <|> return () :: AStar () () () ())
            `shouldBe`
              (Nothing :: Maybe ((), ()))
    -- describe "tryWhile" $ do
    --     it "should stop if weight gets too high" $ do
    --           -- Use tuple monad to see how far we get
    --           do flip (tryWhileT (const (< 4))) () $ do
    --                 asum [ estimate (10 :: Sum Int) >> lift ([10], ()) >> empty
    --                       , estimate (1 :: Sum Int) >> lift ([1], ()) >> empty
    --                       , estimate (5 :: Sum Int) >> lift ([5], ()) >> empty
    --                       , estimate (3 :: Sum Int) >> lift ([3], ()) >> empty
    --                      ]
    --         `shouldBe`
    --           ([1, 3] :: [Int], Nothing :: Maybe ((), ()))


distanceTo :: (Int, Int) -> (Int, Int) -> Int
distanceTo (x, y) (x', y') = abs (x - x') + abs (y - y')

move :: Monad m => Move -> AStarT Context (Sum Int) r m ()
move L = moves <>= [L] >> current . _1 -= 1
move R = moves <>= [R] >> current . _1 += 1
move U = moves <>= [U] >> current . _2 -= 1
move D = moves <>= [D] >> current . _2 += 1

findPoint :: AStar Context (Sum Int) (Int, Int) ()
findPoint = do
    c <- use current
    gl <- use goal
    when (c == gl) $ done gl
    estimate . Sum $ distanceTo gl c
    asum
        [ move L >> findPoint
        , move R >> findPoint
        , move U >> findPoint
        , move D >> findPoint
        ]

-- To make things more interesting we'll say that moving through coordinates
-- where either 'x' or 'y' are divisible by the other is 3X more expensive!
addCost :: MonadAStar (Sum Int) r m => (Int, Int) -> m ()
-- Don't divide by zero!
addCost (0, _) = spend 1
addCost (_, 0) = spend 1
addCost (x, y)
    | mod x y == 0 || mod y x == 0 = spend 3
addCost _ = spend 1


mineField :: AStar Context (Sum Int) (Int, Int) ()
mineField = do
    -- Get the current position
    (x, y) <- use current
    -- Add the current location to our list
    locations <>= [(x, y)]
    -- If we're at the goal we're done!
    gl <- use goal
    when ((x, y) == gl)
      $ done gl
    -- Add the cost of the current position
    addCost (x, y)
    -- Estimate steps to completion
    estimate . Sum $ distanceTo gl (x, y)
    asum
        [ move L >> mineField
        , move R >> mineField
        , move U >> mineField
        , move D >> mineField
        ]