PieceMap-based dispatch

yesod-routes/Yesod/Routes.hs

@@ -9,8 +9,11 @@ import Data.Text (Text)
 import Web.ClientSession (Key)
 import Yesod.Core (Route)
 import qualified Data.Vector as V
-import Data.Maybe (fromMaybe)
+import Data.Maybe (fromMaybe, mapMaybe, listToMaybe)
 import qualified Data.Map as Map
+import Data.List (sortBy)
+import Data.Ord (comparing)
+import Control.Arrow (second)
 
 data Piece = StaticPiece Text | SinglePiece
 
@@ -30,39 +33,75 @@ toDispatch rhs =
 
 bcToDispatch :: ByCount sub master res -> Dispatch sub master res
 bcToDispatch (ByCount vec rest) sub mkey ts master toMaster =
-    go rhs
+    case go ts rhs of
+        Nothing -> Nothing
+        Just dispatch -> dispatch sub mkey ts master toMaster
   where
     len = length ts
     rhs = fromMaybe rest $ vec V.!? len
 
-    go [] = Nothing
-    go (x:xs) = maybe (go xs) Just $ if checkStatics ts (rhPieces x) (rhHasMulti x) then rhHandler x sub mkey ts master toMaster else Nothing
+go :: [Text]
+   -> PieceMap sub master res
+   -> Maybe (Dispatch sub master res)
+go _ (PieceMapEnd r) =
+    listToMaybe $ map snd $ sortBy (comparing fst) r
+go (t:ts) (PieceMap dyn sta) = go ts $
+    case Map.lookup t sta of
+        Nothing -> dyn
+        Just pm -> append dyn pm
+go [] _ = Nothing
 
-    checkStatics [] [] _ = True
-    checkStatics [] _ _ = False
-    checkStatics _ [] isMulti = isMulti
-    checkStatics (_:paths) (SinglePiece:pieces) isMulti = checkStatics paths pieces isMulti
-    checkStatics (path:paths) (StaticPiece piece:pieces) isMulti =
-        path == piece && checkStatics paths pieces isMulti
+append :: PieceMap a b c -> PieceMap a b c -> PieceMap a b c
+append (PieceMapEnd a) (PieceMapEnd b) = PieceMapEnd $ a ++ b
+append (PieceMap a x) (PieceMap b y) =
+    PieceMap (append a b) (Map.unionWith append x y)
+append _ _ = error "Mismatched PieceMaps for append"
 
 data PieceMap sub master res = PieceMap
-    { pmHandlers :: Either (PieceMap sub master res) [(Int, RouteHandler sub master res)]
+    { pmDynamic :: PieceMap sub master res
     , pmStatic :: Map.Map Text (PieceMap sub master res)
-    }
+    } | PieceMapEnd [(Int, Dispatch sub master res)]
 
-toPieceMap :: [RouteHandler sub master res] -> PieceMap sub master res
-toPieceMap = undefined
+toPieceMap :: Int -> [RouteHandler sub master res] -> PieceMap sub master res
+toPieceMap depth = toPieceMap' depth . zip [1..]
+
+toPieceMap' :: Int
+            -> [(Int, RouteHandler sub master res)]
+            -> PieceMap sub master res
+toPieceMap' 0 rhs =
+    PieceMapEnd $ take 1
+                $ map (second rhHandler)
+                $ sortBy (comparing fst) rhs
+toPieceMap' depth rhs = PieceMap
+    { pmDynamic = toPieceMap' depth' dynamics
+    , pmStatic = Map.map (toPieceMap' depth') statics
+    }
+  where
+    depth' = depth - 1
+
+    pairs = map toPair rhs
+    toPair (i, RouteHandler (p:ps) b c) = (p, (i, RouteHandler ps b c))
+    toPair _ = error "toPieceMap' received a route with empty pieces"
+
+    getDynamic (SinglePiece, rh) = Just rh
+    getDynamic _ = Nothing
+    dynamics = mapMaybe getDynamic pairs
+
+    getStatic (StaticPiece t, rh) = Just $ Map.singleton t [rh]
+    getStatic _ = Nothing
+    statics = Map.unionsWith (++) $ mapMaybe getStatic pairs
 
 data ByCount sub master res = ByCount
-    { bcVector :: !(V.Vector [RouteHandler sub master res])
-    , bcRest :: ![RouteHandler sub master res]
+    { bcVector :: !(V.Vector (PieceMap sub master res))
+    , bcRest :: !(PieceMap sub master res)
     }
 
 toBC :: [RouteHandler sub master res] -> ByCount sub master res
 toBC rhs =
     ByCount
-        { bcVector = V.map (\i -> filter (canHaveLength i) rhs) $ V.enumFromN 0 (maxLen + 1)
-        , bcRest = filter rhHasMulti rhs
+        { bcVector = V.map (\i -> toPieceMap i $ filter (canHaveLength i) rhs)
+                   $ V.enumFromN 0 (maxLen + 1)
+        , bcRest = toPieceMap maxLen $ filter rhHasMulti rhs
         }
   where
     maxLen = maximum $ map (length . rhPieces) rhs