restructure GTFS types

unfortunately doesn't quite get rid of all the type family still
since it's just too useful … but does reduce it somewhat.

Also, maps are much easier for looking things up than vectors!

2 files changed, 81 insertions, 60 deletions

diff --git a/lib/GTFS.hs b/lib/GTFS.hs
index 35a85ea..9ad01f1 100644
--- a/lib/GTFS.hs
+++ b/lib/GTFS.hs
@@ -13,6 +13,11 @@
 {-# LANGUAGE GeneralizedNewtypeDeriving #-}
 {-# LANGUAGE FlexibleContexts #-}
 {-# LANGUAGE UndecidableInstances #-}
+{-# LANGUAGE TupleSections #-}
+{-# LANGUAGE TupleSections #-}
+{-# LANGUAGE TupleSections #-}
+{-# LANGUAGE TupleSections #-}
+{-# LANGUAGE TupleSections #-}
 
 
 module GTFS where
@@ -47,7 +52,9 @@ import GHC.Generics (Generic)
 -- import Data.Aeson.Generic (Options(fieldLabelModifier), deriveJSON, defaultOptions)
 import qualified Data.Text as T
 import Data.Char (toLower)
-
+import Data.Map (Map)
+import qualified Data.Map as M
+import Data.Foldable (Foldable(fold))
 
 aesonOptions prefix =
   defaultOptions { fieldLabelModifier = fieldModifier (T.length prefix) }
@@ -119,7 +126,7 @@ data Stop (deep :: Depth) = Stop
   { stopTrip :: TripID
   , stopArrival :: Time
   , stopDeparture :: Time
-  , stopStation:: Switch deep Station StationID
+  , stopStation :: Switch deep Station StationID
   , stopSequence :: Int
   } deriving Generic
 
@@ -253,31 +260,39 @@ instance CSV.FromNamedRecord (Trip Shallow) where
     <*> r .: "shape_id"
     <*> pure ()
 
-data GTFS (depth :: Depth) = GTFS
-  { stations :: Vector Station
-  , stops :: Vector (Stop depth)
-  , trips :: Vector (Trip depth)
-  , calendar :: Maybe (Vector Calendar)
-  , calendarDates :: Maybe (Vector CalendarDate)
+data RawGTFS = RawGTFS
+  { rawStations :: Vector Station
+  , rawStops :: Vector (Stop Shallow)
+  , rawTrips :: Vector (Trip Shallow)
+  , rawCalendar :: Maybe (Vector Calendar)
+  , rawCalendarDates :: Maybe (Vector CalendarDate)
   }
 
-deriving instance Show (GTFS Shallow)
-deriving instance Show (GTFS Deep)
 
-class Loadable depth where
-  loadGtfs :: FilePath -> IO (GTFS depth)
+data GTFS = GTFS
+  { stations :: Map StationID Station
+  , trips :: Map TripID (Trip Deep)
+  , calendar :: Map DayOfWeek (Vector Calendar)
+  , calendarDates :: Map Day (Vector CalendarDate)
+
+  , fancyCalendar :: Day -> (Vector ServiceID, Vector (Trip Deep))
+  -- ^ a more "fancy" encoding of the calendar?
+  } -- deriving Show
+
+
 
-instance Loadable Shallow where
-  loadGtfs path = do
+loadRawGtfs :: FilePath -> IO RawGTFS
+loadRawGtfs path = do
     zip <- Zip.toArchive <$> LB.readFile "./gtfs.zip"
-    GTFS <$> decodeTable' "stops.txt" zip
-         <*> decodeTable' "stop_times.txt" zip
-         <*> decodeTable' "trips.txt" zip
-         <*> decodeTable "calendar.txt" zip
-         <*> decodeTable "calendar_dates.txt" zip
+    RawGTFS
+      <$> decodeTable' "stops.txt" zip
+      <*> decodeTable' "stop_times.txt" zip
+      <*> decodeTable' "trips.txt" zip
+      <*> decodeTable "calendar.txt" zip
+      <*> decodeTable "calendar_dates.txt" zip
     where
       decodeTable :: CSV.FromNamedRecord a => FilePath -> Zip.Archive -> IO (Maybe (Vector a))
-      decodeTable path zip = do
+      decodeTable path zip =
         case Zip.findEntryByPath path zip of
           Nothing -> pure Nothing
           Just csv -> case CSV.decodeByName (Zip.fromEntry csv) of
@@ -288,13 +303,37 @@ instance Loadable Shallow where
           Nothing -> fail $ "required file "+|path|+" not found in gtfs.zip"
           Just a -> pure a
 
-instance Loadable Deep where
-  loadGtfs path = do
-    shallow <- loadGtfs @Shallow path
-    stops' <- V.mapM (pushStop (stations shallow)) (stops shallow)
-    trips' <- V.mapM (pushTrip stops') (trips shallow)
-    pure $ shallow { stops = stops', trips = trips' }
+loadGtfs :: FilePath -> IO GTFS
+loadGtfs path = do
+    shallow@RawGTFS{..} <- loadRawGtfs path
+    stops' <- V.mapM (pushStop rawStations) rawStops
+    trips' <- V.mapM (pushTrip stops') rawTrips
+    pure $ GTFS
+      { stations =
+        M.fromList $ (\station -> (stationId station, station))
+        <$> V.toList rawStations
+      , trips =
+        M.fromList $ (\trip -> (tripTripID trip, trip))
+        <$> V.toList trips'
+      , calendar =
+        fmap V.fromList
+        $ M.fromListWith (<>)
+        $ concatMap (\cal -> (, [cal]) <$> weekdays cal)
+        $ V.toList (fromMaybe mempty rawCalendar)
+      , calendarDates =
+        fmap V.fromList
+        $ M.fromListWith (<>) $ (\cd -> (caldateDate cd, [cd]))
+        <$> V.toList (fromMaybe mempty rawCalendarDates)
+      }
     where
+      weekdays Calendar{..} =
+        if calMonday then [Monday] else []
+        <> if calTuesday then [Tuesday] else []
+        <> if calWednesday then [Wednesday] else []
+        <> if calThursday then [Thursday] else []
+        <> if calFriday then [Friday] else []
+        <> if calSaturday then [Saturday] else []
+        <> [Sunday | calSunday]
       pushStop :: Vector Station -> Stop Shallow -> IO (Stop Deep)
       pushStop stations stop = do
         station <- case tableLookup stationId (stopStation stop) stations of
@@ -310,41 +349,25 @@ instance Loadable Deep where
 
 
 
-servicesOnDay :: GTFS Deep -> Day -> Vector ServiceID
+servicesOnDay :: GTFS -> Day -> Vector ServiceID
 servicesOnDay GTFS{..} day =
   fmap caldateServiceId added <> V.filter notCancelled regular
-  where (added,removed) = case calendarDates of
-          Nothing -> (mempty,mempty)
-          Just exs ->
-            V.partition (\cd -> caldateExceptionType cd == ServiceAdded)
-            $ V.filter (\cd -> caldateDate cd == day) exs
-        regular = case calendar of
-          Nothing -> mempty
-          Just cs -> V.mapMaybe (\cal -> if isRunning cal then Just (calServiceId cal) else Nothing) cs
-          where isRunning Calendar{..} =
-                  day >= calStartDate &&
-                  day <= calEndDate &&
-                  case weekday of
-                    Monday -> calMonday
-                    Tuesday -> calTuesday
-                    Wednesday -> calWednesday
-                    Thursday -> calThursday
-                    Friday -> calFriday
-                    Saturday -> calSaturday
-                    Sunday -> calSunday
-        weekday = dayOfWeek day
+  where (added,removed) =
+          V.partition (\cd -> caldateExceptionType cd == ServiceAdded)
+          . fromMaybe mempty $ M.lookup day calendarDates
+        regular = maybe mempty (fmap calServiceId) $ M.lookup (dayOfWeek day) calendar
         notCancelled serviceID =
           null (tableLookup caldateServiceId serviceID removed)
 
-tripsOfService :: GTFS Deep -> ServiceID -> Vector (Trip Deep)
+tripsOfService :: GTFS -> ServiceID -> Map TripID (Trip Deep)
 tripsOfService GTFS{..} serviceId =
-  V.filter (\trip -> tripServiceId trip == serviceId ) trips
+  M.filter (\trip -> tripServiceId trip == serviceId ) trips
 
 -- TODO: this should filter out trips ending there
-tripsAtStation :: GTFS Deep -> StationID -> Vector TripID
+tripsAtStation :: GTFS -> StationID -> Vector TripID
 tripsAtStation GTFS{..} at = fmap stopTrip stops
   where
     stops = V.filter (\(stop :: Stop Deep) -> stationId (stopStation stop) == at) stops
 
-tripsOnDay :: GTFS Deep -> Day -> Vector (Trip Deep)
-tripsOnDay gtfs today = V.concatMap (tripsOfService gtfs) (servicesOnDay gtfs today)
+tripsOnDay :: GTFS -> Day -> Map TripID (Trip Deep)
+tripsOnDay gtfs today = foldMap (tripsOfService gtfs) (servicesOnDay gtfs today)
diff --git a/lib/Server.hs b/lib/Server.hs
index 7a79aa8..f9bf36b 100644
--- a/lib/Server.hs
+++ b/lib/Server.hs
@@ -50,8 +50,8 @@ instance ToJSON TrainPing where
 
 type KnownTrips = TVar (Map Token [TrainPing])
 
-type API = "stations" :> Get '[JSON] (Vector Station)
-  :<|> "timetable" :> Capture "Station ID" StationID :> Get '[JSON] (Vector (Trip Deep))
+type API = "stations" :> Get '[JSON] (Map StationID Station)
+  :<|> "timetable" :> Capture "Station ID" StationID :> Get '[JSON] (Map TripID (Trip Deep))
   :<|> "trip" :> Capture "Trip ID" TripID :> Get '[JSON] (Trip Deep)
   -- ingress API (put this behind BasicAuth?)
   -- TODO: perhaps require a first ping for registration?
@@ -61,28 +61,26 @@ type API = "stations" :> Get '[JSON] (Vector Station)
   -- debug things
   :<|> "debug" :> "state" :> Get '[JSON] (Map Token [TrainPing])
 
-server :: GTFS Deep -> KnownTrips -> Server API
+server :: GTFS -> KnownTrips -> Server API
 server gtfs@GTFS{..} knownTrains = handleStations :<|> handleTimetable :<|> handleTrip
   :<|> handleRegister :<|> handleTrainPing :<|> handleDebugState
   where handleStations = pure stations
         handleTimetable station = do
           today <- liftIO getCurrentTime <&> utctDay
           pure $ tripsOnDay gtfs today
-        handleTrip trip = case tableLookup tripTripID trip trips of
+        handleTrip trip = case M.lookup trip trips of
           Just res -> pure res
           Nothing -> throwError err404
         handleRegister tripID = liftIO $ do
           token <- UUID.nextRandom <&> Token
           atomically $ modifyTVar knownTrains (M.insert token [])
           pure token
-        handleTrainPing token ping = liftIO $ do
-          putStrLn "got train ping"
-          atomically $ do
+        handleTrainPing token ping = liftIO $ atomically $ do
             modifyTVar knownTrains (M.update (\history -> Just (ping : history)) token)
             pure ()
         handleDebugState = liftIO $ readTVarIO knownTrains
 
-application :: GTFS Deep -> IO Application
+application :: GTFS -> IO Application
 application gtfs = do
   knownTrips <- newTVarIO mempty
   pure $ serve (Proxy @API) $ server gtfs knownTrips