4 files changed, 31 insertions, 10 deletions

diff --git a/Handler/InventoryListing.hs b/Handler/InventoryListing.hs
index cbf4eab..c2ec5d1 100644
--- a/Handler/InventoryListing.hs
+++ b/Handler/InventoryListing.hs
@@ -16,7 +16,7 @@ postInventoryListingR = do
     FormFailure errors -> mapM_ (addMessage "formError" . toHtml) errors
     _ -> return ()
 
-  (stockSort -> stock) <- runDB $ mapM withType =<< selectList [] []
+  (stockSort -> stock) <- runDB $ withTypes =<< selectList [] []
 
   selectRep $ do
     provideJson (stock :: [WithType (Entity Item)])
diff --git a/Handler/ReferenceListing.hs b/Handler/ReferenceListing.hs
index e433429..690f3f7 100644
--- a/Handler/ReferenceListing.hs
+++ b/Handler/ReferenceListing.hs
@@ -17,7 +17,7 @@ postReferenceListingR = do
     FormFailure errors -> mapM_ (addMessage "formError" . toHtml) errors
     _ -> return ()
 
-  (referenceSort -> reference) <- runDB $ mapM withType =<< selectList [] [Asc ReferenceKind]
+  (referenceSort -> reference) <- runDB $ withTypes =<< selectList [] [Asc ReferenceKind]
 
   selectRep $ do
     provideJson (reference :: [WithType (Entity Reference)])
diff --git a/Handler/UpdateItem.hs b/Handler/UpdateItem.hs
index 503d8ca..b4c8713 100644
--- a/Handler/UpdateItem.hs
+++ b/Handler/UpdateItem.hs
@@ -28,7 +28,7 @@ postUpdateItemR fsUpdateId = do
     provideRep $ case updateResult of
       FormSuccess _ -> redirect $ InventoryListingR :#: fsUpdateId :: Handler Html
       _ -> do
-        (stockSort -> stock) <- runDB $ mapM withType =<< selectList [] []
+        (stockSort -> stock) <- runDB $ withTypes =<< selectList [] []
         defaultLayout $ inventoryListing InventoryState
           { invFormState = Just UpdateForm{..}
           , ..
diff --git a/Model.hs b/Model.hs
index 13f18b8..90f5904 100644
--- a/Model.hs
+++ b/Model.hs
@@ -10,6 +10,9 @@ import Control.Monad.Writer
 import           Data.Text   (Text)
 import qualified Data.Text as Text
 
+import           Data.Map   (Map, (!))
+import qualified Data.Map as Map
+
 import qualified Data.HashMap.Lazy as HashMap
 
 import Data.Aeson
@@ -23,19 +26,24 @@ share [mkPersist sqlSettings, mkMigrate "migrateAll"]
     $(persistFileWith lowerCaseSettings "config/models")
 
 class HasType a where
-  getType :: ( BaseBackend backend ~ SqlBackend
-             , MonadIO m
-             , PersistStoreRead backend
-             ) => a -> ReaderT backend m Kind
+  fkType :: a -> Key Kind
+  
+getType :: ( BaseBackend backend ~ SqlBackend
+           , MonadIO m
+           , PersistStoreRead backend
+           , HasType a
+           ) => a -> ReaderT backend m Kind
+getType = getJust . fkType
 
 instance HasType Item where
-  getType = belongsToJust itemFkType
+  fkType = itemFkType
 
 instance HasType Reference where
-  getType = belongsToJust referenceFkType
+  fkType = referenceFkType
 
 instance HasType a => HasType (Entity a) where
-  getType Entity{..} = getType entityVal
+  fkType = fkType . entityVal
+    
 
 withType :: ( BaseBackend backend ~ SqlBackend
             , MonadIO m
@@ -44,6 +52,19 @@ withType :: ( BaseBackend backend ~ SqlBackend
             ) => a -> ReaderT backend m (WithType a)
 withType val = (val `WithType`) . kindType <$> getType val
 
+withTypes :: ( BaseBackend backend ~ SqlBackend
+             , MonadIO m
+             , PersistStoreRead backend
+             , PersistQueryRead backend
+             , HasType a
+             , Functor f
+             ) => f a -> ReaderT backend m (f (WithType a))
+withTypes vals = do
+  typeMap <- Map.fromList . map (\(Entity kindKey kind) -> (kindKey, kind)) <$> selectList [] []
+  let
+    lookupType x = (x `WithType`) . kindType $ typeMap ! fkType x
+  return $ lookupType <$> vals
+
 instance Ord Item where
   x `compare` y = mconcat cmprs
     where