matt-noonan / justified-containers

Standard containers, with keys that carry type-level proofs of their own presence.
https://hackage.haskell.org/package/justified-containers
BSD 2-Clause "Simplified" License
80 stars 5 forks source link

Speed up withRecMap #12

Open treeowl opened 1 year ago

treeowl commented 1 year ago

Currently, withRecMap rebuilds the map from scratch (not even recognizing that the keys are strictly ascending), which is a wasteful $O(n \log n)$. Even realizing the ok list is fairly bad, especially if the map is large. I would probably do this:

withRecMap
  :: (Ord k, Traversable f)
  => M.Map k (f k)
  -> (forall ph. Map ph k (f (Key ph k)) -> t)
  -> Either [MissingReference k f] t
withRecMap (Map m) cont =
  case bad of
    [] -> Right $ cont (Map $ M.map (fmap Key) $ m)
    _  -> Left (map (\(k,v) -> (k, fmap (id &&& locate) v)) bad)
  where
    bad = filter (any ((== Missing) . locate)) (M.elems m)
    locate k = if M.member k m then Present else Missing
{-# INLINABLE withRecMap #-}

In the (moderately common) case where the functor has an fmap/coerce rule, or simple enough that GHC can work out that fmap Key = coerce, M.map (fmap Key) should end up returning the original map intact. In other cases, rebuilding will be $O(n)$, which is an improvement.

treeowl commented 1 year ago

Oh, I see a different change has been made in this direction in the repo, but not on Hackage. Interesting....