16-1: more mutable data structures for another 7x speedup

1 files changed, 40 insertions, 37 deletions

diff --git a/16-1.hs b/16-1.hs
index 69eb0cd..884f5ec 100644
--- a/16-1.hs
+++ b/16-1.hs
@@ -1,18 +1,21 @@
-import Control.Monad (foldM)
+import Control.Monad (filterM, foldM)
 import Control.Monad.ST
 import Control.Monad.Trans.Class
 import Control.Monad.Trans.Maybe
 import Data.Array as A
+import Data.Array.ST
 import Data.Bifunctor (second)
 import Data.List as L (find, uncons)
 import Data.Maybe (fromJust)
 import Data.Tuple as T (swap)
-import Data.Vector.Mutable as V hiding (foldM)
+import Data.Vector.Mutable as V hiding (foldM, mapM_)
 import Prelude
 import Prelude as P (length)
 
 data BucketQueue s a = BQ Int (V.MVector s [a])
 
+instance Show (BucketQueue s a) where show (BQ top _) = show top
+
 insert :: BucketQueue s a -> (Int,a) -> ST s (BucketQueue s a)
 insert (BQ top buckets) (prio,x) =
   V.modify buckets (x:) prio >> return (BQ (min top prio) buckets)
@@ -63,7 +66,7 @@ data Tile = Wall | Tile { untile :: [Direction] }
 
 type Map = A.Array (Int,Int) Tile
 
-type Costs = A.Array Node (Maybe Int)
+type Costs s = STArray s Node (Maybe Int)
 
 neighbours :: (Int,Int) -> [(Direction, (Int,Int))]
 neighbours (x,y) = [ (North, (x,y-1))
@@ -126,59 +129,59 @@ betterPath m through dest@(node,oldCost) =
      then n
      else dest
 
-cost :: Costs -> Node -> (Node, Maybe Int)
-cost costs node = (node, costs A.! node)
+cost :: Costs s -> Node -> ST s (Node, Maybe Int)
+cost costs node = (node,) <$> readArray costs node
 
-data PathingResult s = Success Int | Fail | Working (Frontier s, Costs)
+data PathingResult s = Success Int | Fail | Working (Frontier s, Costs s)
 
 heuristic :: (Int,Int) -> (Node,Int) -> (Node,Int)
 heuristic goal (n@(p,_),f) = (n, f + distance goal p)
 
+(///) :: (MArray a e m, Ix i, Show i, Show e) => a i e -> [(i,e)] -> m ()
+a /// asss = mapM_ (uncurry (writeArray a)) asss
+
+unvisited :: Costs s -> Node -> ST s Bool
+unvisited costs node = (== Nothing) <$> readArray costs node
+
 pathOnce :: Map
          -> (Int,Int)
          -> Frontier s
-         -> Costs
+         -> Costs s
          -> ST s (PathingResult s)
 pathOnce m goal frontier costs =
   runMaybeT (view frontier) >>= maybe (return Fail) continue
   where
-    continue :: (Node,Frontier s) -> ST s (PathingResult s)
     continue (current@(p,_), rest)
-      | p == goal = return . Success $ fromJust (costs A.! current)
-      | otherwise =
-          let costs' = {-# SCC "costs" #-}
-                costs A.// [ betterPath m (cost costs current) n
-                           | n <- expand m (cost costs current)
-                           ]
-              frontier' = foldM insert rest
-                          . map (T.swap . heuristic goal . second fromJust)
-                          . filter ((== Nothing) . (costs A.!) . fst)
-                          $ expand m (cost costs current)
-          in Working . (, costs') <$> frontier'
+      | p == goal = Success . fromJust <$> readArray costs current
+      | otherwise = do
+          curCost <- cost costs current
+          let nexts = expand m curCost
+          frontier' <- filterM (unvisited costs . fst) nexts
+            >>= foldM insert rest
+            . map (T.swap . heuristic goal . second fromJust)
+          costs /// map (betterPath m curCost) nexts
+          return . Working $ (frontier', costs)
 
 findPath :: Map -> (Int,Int) -> (Int,Int) -> ST s Int
-findPath m start goal = frontier >>= flip loop costs
-  where
-    frontier = {-# SCC "frontier" #-}
+findPath m start goal = do
+    frontier <- {-# SCC "frontier" #-}
       let (_,(w,h)) = bounds m
       in newBQ (100*w*h) >>= flip insert (0,(start,East))
 
-    costs = let (lower,upper) = bounds m
-                in array ((lower,North),(upper,West))
-                   [ if p == start && d == East
-                     then ((p,d),Just 0)
-                     else ((p,d),Nothing)
-                   | p <- indices m
-                   , d <- [North, South, East, West]
-                   ]
+    costs <- let (lower,upper) = bounds m
+             in newArray ((lower,North),(upper,West)) Nothing
+
+    writeArray costs (start,East) (Just 0)
 
-    loop :: Frontier s -> Costs -> ST s Int
-    loop frontier costs = do
-      once <- pathOnce m goal frontier costs
-      case once of
-        Success x -> return x
-        Fail -> error "no path"
-        Working (frontier', costs') -> loop frontier' costs'
+    loop frontier costs
+      where
+        loop :: Frontier s -> Costs s -> ST s Int
+        loop frontier costs = do
+          once <- pathOnce m goal frontier costs
+          case once of
+            Success x -> return x
+            Fail -> error "no path"
+            Working (frontier', costs') -> loop frontier' costs'
 
 main :: IO ()
 main = do