16-1: mutable vectors for 42-fold speedup

1 files changed, 70 insertions, 49 deletions

diff --git a/16-1.hs b/16-1.hs
index a8640fb..69eb0cd 100644
--- a/16-1.hs
+++ b/16-1.hs
@@ -1,33 +1,39 @@
+import Control.Monad (foldM)
+import Control.Monad.ST
+import Control.Monad.Trans.Class
+import Control.Monad.Trans.Maybe
 import Data.Array as A
 import Data.Bifunctor (second)
 import Data.List as L (find, uncons)
 import Data.Maybe (fromJust)
-import Data.Tuple (swap)
-import Data.Vector as V hiding
-  ((++), elem, filter, foldl, head, length, map, zip)
-import Prelude hiding (replicate)
+import Data.Tuple as T (swap)
+import Data.Vector.Mutable as V hiding (foldM)
+import Prelude
+import Prelude as P (length)
 
-data BucketQueue a = BQ Int (V.Vector [a]) deriving Show
+data BucketQueue s a = BQ Int (V.MVector s [a])
 
-insert :: BucketQueue a -> (Int,a) -> BucketQueue a
+insert :: BucketQueue s a -> (Int,a) -> ST s (BucketQueue s a)
 insert (BQ top buckets) (prio,x) =
-  let top' = min top prio
-  in BQ top' $ V.accum (flip (:)) buckets [ (prio, x) ]
+  V.modify buckets (x:) prio >> return (BQ (min top prio) buckets)
 
-view :: BucketQueue a -> Maybe (a, BucketQueue a)
+view :: BucketQueue s a -> MaybeT (ST s) (a, BucketQueue s a)
 view (BQ top buckets) = do
-  (x, xs) <- buckets V.!? top >>= L.uncons
-  let buckets' = buckets V.// [ (top, xs) ]
-  let top' = findTop top buckets'
-  return (x, BQ top' buckets')
-
-findTop :: Int -> V.Vector [a] -> Int
-findTop n v = case v V.! n of
-  [] -> if n < length v - 1 then findTop (n+1) v else n
-  _  -> n
-
-newBQ :: Int -> BucketQueue a
-newBQ n = BQ n $ replicate n []
+  (x, xs) <- V.read buckets top >>= hoistMaybe . L.uncons
+  lift $ V.write buckets top xs
+  top' <- lift $ findTop top buckets
+  return (x, BQ top' buckets)
+{-# SCC view #-}
+
+findTop :: Int -> V.MVector s [a] -> ST s Int
+findTop n v = do
+  top <- V.read v n
+  case top of
+    [] -> if n < V.length v - 1 then findTop (n+1) v else return n
+    _  -> return n
+
+newBQ :: Int -> ST s (BucketQueue s a)
+newBQ n = BQ n <$> V.replicate n []
 
 next :: Direction -> (Int,Int) -> (Int,Int)
 next North (x,y) = (x,y-1)
@@ -49,7 +55,7 @@ right West  = North
 
 type Node = ((Int,Int), Direction)
 
-type Frontier = BucketQueue Node
+type Frontier s = BucketQueue s Node
 
 data Direction = North | South | East | West deriving (Eq,Ix,Ord,Show)
 
@@ -68,8 +74,8 @@ neighbours (x,y) = [ (North, (x,y-1))
 
 parseMap :: [String] -> Map
 parseMap s =
-  let w = length (head s) - 1
-      h = length s - 1
+  let w = P.length (head s) - 1
+      h = P.length s - 1
       f p = case g p of
         '#' -> Wall
         _   -> Tile . map fst . filter ((/= '#') . g . snd) $ neighbours p
@@ -82,16 +88,16 @@ parseMap s =
 findStart :: [String] -> (Int,Int)
 findStart ls =
   let Just p = L.find ((== 'S') . snd) [ ((x,y), ls !! y !! x)
-                                       | x <- [0..length (head ls) - 1]
-                                       , y <- [0..length ls - 1]
+                                       | x <- [0..P.length (head ls) - 1]
+                                       , y <- [0..P.length ls - 1]
                                        ]
   in fst p
 
 findGoal :: [String] -> (Int,Int)
 findGoal ls =
   let Just p = L.find ((== 'E') . snd) [ ((x,y), ls !! y !! x)
-                                       | x <- [0..length (head ls) - 1]
-                                       , y <- [0..length ls - 1]
+                                       | x <- [0..P.length (head ls) - 1]
+                                       , y <- [0..P.length ls - 1]
                                        ]
   in fst p
 
@@ -107,9 +113,13 @@ expand m ((pos, dir), Just cost) =
   ++ ((\x -> ((pos,x), Just $ cost + 1000)) <$> [left dir, right dir])
 
 pathThrough :: Map -> (Node,Maybe Int) -> Node -> (Node,Maybe Int)
-pathThrough m through dest = head . filter ((== dest) . fst) $ expand m through
+pathThrough m through dest =
+  head . filter ((== dest) . fst) $ expand m through
 
-betterPath :: Map -> (Node,Maybe Int) -> (Node,Maybe Int) -> (Node,Maybe Int)
+betterPath :: Map
+           -> (Node,Maybe Int)
+           -> (Node,Maybe Int)
+           -> (Node,Maybe Int)
 betterPath m through dest@(node,oldCost) =
   let n@(_,newCost) = pathThrough m through node
   in if newCost < oldCost
@@ -119,30 +129,39 @@ betterPath m through dest@(node,oldCost) =
 cost :: Costs -> Node -> (Node, Maybe Int)
 cost costs node = (node, costs A.! node)
 
-data PathingResult = Success Int | Fail | Working (Frontier,Costs)
+data PathingResult s = Success Int | Fail | Working (Frontier s, Costs)
 
 heuristic :: (Int,Int) -> (Node,Int) -> (Node,Int)
 heuristic goal (n@(p,_),f) = (n, f + distance goal p)
 
-pathOnce :: Map -> (Int,Int) -> Frontier -> Costs -> PathingResult
-pathOnce m goal frontier costs = maybe Fail continue (view frontier)
+pathOnce :: Map
+         -> (Int,Int)
+         -> Frontier s
+         -> Costs
+         -> 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 = Success $ fromJust (costs A.! current)
+      | p == goal = return . Success $ fromJust (costs A.! current)
       | otherwise =
-          let costs' = costs A.// [ betterPath m (cost costs current) n
-                                  | n <- expand m (cost costs current)
-                                  ]
-              frontier' = foldl insert rest
-                          . map (swap . heuristic goal . second fromJust)
+          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 (frontier', costs')
+          in Working . (, costs') <$> frontier'
 
-findPath :: Map -> (Int,Int) -> (Int,Int) -> Int
-findPath m start goal = loop frontier costs
+findPath :: Map -> (Int,Int) -> (Int,Int) -> ST s Int
+findPath m start goal = frontier >>= flip loop costs
   where
-    frontier = insert (newBQ $ 10^6) (0,(start,East))
+    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))
@@ -153,11 +172,13 @@ findPath m start goal = loop frontier costs
                    , d <- [North, South, East, West]
                    ]
 
-    loop :: Frontier -> Costs -> Int
-    loop frontier costs = case pathOnce m goal frontier costs of
-      Success x -> x
-      Fail -> error "no path"
-      Working (frontier', costs') -> loop frontier' costs'
+    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'
 
 main :: IO ()
 main = do
@@ -165,4 +186,4 @@ main = do
   let m = parseMap input
   let start = findStart input
   let goal = findGoal input
-  print $ findPath m start goal
+  print (runST $ findPath m start goal :: Int)