From fbf58cbc95e065084123feec0adc5bb23427cb38 Mon Sep 17 00:00:00 2001
From: Michael Snoyman <michael@snoyman.com>
Date: Fri, 7 Oct 2011 08:22:27 +0200
Subject: [PATCH] Better overlapping support

---
 yesod-routes/Yesod/Routes.hs | 36 +++++++++++++++++++++---------------
 yesod-routes/test/main.hs    | 12 ++++++++++++
 2 files changed, 33 insertions(+), 15 deletions(-)

diff --git a/yesod-routes/Yesod/Routes.hs b/yesod-routes/Yesod/Routes.hs
index e0d83318..4ab1a906 100644
--- a/yesod-routes/Yesod/Routes.hs
+++ b/yesod-routes/Yesod/Routes.hs
@@ -9,7 +9,7 @@ import Data.Text (Text)
 import Web.ClientSession (Key)
 import Yesod.Core (Route)
 import qualified Data.Vector as V
-import Data.Maybe (fromMaybe, mapMaybe, listToMaybe)
+import Data.Maybe (fromMaybe, mapMaybe)
 import qualified Data.Map as Map
 import Data.List (sortBy)
 import Data.Ord (comparing)
@@ -33,28 +33,33 @@ toDispatch rhs =
 
 bcToDispatch :: ByCount sub master res -> Dispatch sub master res
 bcToDispatch (ByCount vec rest) sub mkey ts master toMaster =
-    case go ts rhs of
-        Nothing -> Nothing
-        Just dispatch -> dispatch sub mkey ts master toMaster
+    go (\x -> x sub mkey ts master toMaster) ts pm
   where
-    len = length ts
-    rhs = fromMaybe rest $ vec V.!? len
+    --pm :: PieceMap sub master res
+    pm = fromMaybe rest $ vec V.!? length ts
 
-go :: [Text]
+go :: (Dispatch sub master res -> Maybe res)
+   -> [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 $
+   -> Maybe res
+go runDispatch _ (PieceMapEnd r) =
+    firstJust runDispatch $ map snd $ sortBy (comparing fst) r
+go runDispatch (t:ts) (PieceMap dyn sta) = go runDispatch ts $
     case Map.lookup t sta of
         Nothing -> dyn
         Just pm -> append dyn pm
-go [] _ = Nothing
+go _ [] _ = Nothing
+
+firstJust :: (a -> Maybe b) -> [a] -> Maybe b
+firstJust _ [] = Nothing
+firstJust f (a:as) = maybe (firstJust f as) Just $ f a
 
 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)
+-- I'm sure there's some nice type-level trickery we could employ here somehow
+-- to ensure this never happens.
 append _ _ = error "Mismatched PieceMaps for append"
 
 data PieceMap sub master res = PieceMap
@@ -69,8 +74,7 @@ toPieceMap' :: Int
             -> [(Int, RouteHandler sub master res)]
             -> PieceMap sub master res
 toPieceMap' 0 rhs =
-    PieceMapEnd $ take 1
-                $ map (second rhHandler)
+    PieceMapEnd $ map (second rhHandler)
                 $ sortBy (comparing fst) rhs
 toPieceMap' depth rhs = PieceMap
     { pmDynamic = toPieceMap' depth' dynamics
@@ -81,7 +85,9 @@ toPieceMap' depth rhs = PieceMap
 
     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"
+    -- if we have no more pieces, that means this is a rhHasMulti, so fill in
+    -- with dynamic
+    toPair (i, RouteHandler [] b c) = (SinglePiece, (i, RouteHandler [] b c))
 
     getDynamic (SinglePiece, rh) = Just rh
     getDynamic _ = Nothing
diff --git a/yesod-routes/test/main.hs b/yesod-routes/test/main.hs
index 51415e2f..7abce23b 100644
--- a/yesod-routes/test/main.hs
+++ b/yesod-routes/test/main.hs
@@ -39,6 +39,13 @@ dynamic = toDispatch
             _ -> error $ "Called dynamic with: " ++ show ts
     ]
 
+overlap :: Dispatch Dummy Dummy Int
+overlap = toDispatch
+    [ RouteHandler [StaticPiece "foo"] False $ result $ const $ Just 20
+    , RouteHandler [StaticPiece "foo"] True $ result $ const $ Just 21
+    , RouteHandler [] True $ result $ const $ Just 22
+    ]
+
 test :: Dispatch Dummy Dummy Int -> [Text] -> Maybe Int
 test dispatch ts = dispatch Dummy Nothing ts Dummy id
 
@@ -65,3 +72,8 @@ main = hspecX $ do
         it "fails correctly on five" $ test dynamic ["five"] @?= Nothing
         it "fails correctly on too many" $ test dynamic ["foo", "baz"] @?= Nothing
         it "fails correctly on too few" $ test dynamic [] @?= Nothing
+    describe "overlap" $ do
+        it "dispatches correctly to foo" $ test overlap ["foo"] @?= Just 20
+        it "dispatches correctly to foo/bar" $ test overlap ["foo", "bar"] @?= Just 21
+        it "dispatches correctly to bar" $ test overlap ["bar"] @?= Just 22
+        it "dispatches correctly to []" $ test overlap [] @?= Just 22