5

I have a computation where I'm inserting values into a Map and then looking them up again. I know that I never use a key before inserting it, but using (!) freely makes me nervous anyway. I'm looking for a way to get a total lookup function that doesn't return a Maybe, and which the type system prevents me from accidentally abusing.

My first thought was to make a monad transformer similar to StateT, where the state is a Map and there are special functions for inserts and lookup in the monad. The insert function returns a Receipt s k newtype, where s is a phantom index type in the style of the ST monad and k is the type of the key, and the lookup function takes a Receipt instead of a bare key. By hiding the Receipt constructor and using a quantified run function similar to runST, this should ensure that lookups only happen after inserts in the same map. (Full code is below.)

But I fear that I've reinvented a wheel, or that that there's an alternate way to get safe, total map lookups that's already in use. Is there any prior art for this problem in a public package somewhere?

{-# LANGUAGE DeriveFunctor, LambdaCase, RankNTypes #-}

module KeyedStateT (KeyedStateT, Receipt, insert, lookup, receiptToKey, runKeyedStateT)
where

import Prelude hiding (lookup)
import Control.Arrow ((&&&))
import Control.Monad (ap, (>=>))
import Data.Map (Map)
import qualified Data.Map as Map
import Data.Maybe (fromJust)

newtype KeyedStateT s k v m a = KeyedStateT (Map k v -> m (a, Map k v)) deriving Functor

keyedState :: Applicative m => (Map k v -> (a, Map k v)) -> KeyedStateT s k v m a
keyedState f = KeyedStateT (pure . f)

instance Monad m => Applicative (KeyedStateT s k v m) where
  pure = keyedState . (,)
  (<*>) = ap

instance Monad m => Monad (KeyedStateT s k v m) where
  KeyedStateT m >>= f = KeyedStateT $ m >=> uncurry ((\(KeyedStateT m') -> m') . f)

newtype Receipt s k = Receipt { receiptToKey :: k }

insert :: (Applicative m, Ord k) => k -> v -> KeyedStateT s k v m (Receipt s k)
insert k v = keyedState $ const (Receipt k) &&& Map.insert k v

lookup :: (Applicative m, Ord k) => Receipt s k -> KeyedStateT s k v m v
lookup (Receipt k) = keyedState $ (Map.! k) &&& id

runKeyedStateT :: (forall s. KeyedStateT s k v m a) -> m (a, Map k v)
runKeyedStateT (KeyedStateT m) = m Map.empty

module Main where

import Data.Functor.Identity (runIdentity)
import qualified KeyedStateT as KS

main = putStrLn . fst . runIdentity $ KS.runKeyedStateT $ do
  one <- KS.insert 1 "hello"
  two <- KS.insert 2 " world"
  h <- KS.lookup one
  w <- KS.lookup two
  pure $ h ++ w

Edit: Several commenters have asked why I want to hold on to a Receipt instead of the actual value. I want to be able to use the Receipt in Sets and Maps (I didn't add the Eq and Ord instances for Receipt in my MVCE, but I have them in my project), but the values in my Map are not equatable. If I replaced Receipt with a key-value pair newtype, I'd have to implement a dishonest Eq instance for that pair that disregarded the value, and then I'd be nervous about that. The Map is there to ensure that there's only one value under consideration for any of my equatable "proxy" keys at any given time.

I suppose an alternate solution that would work just fine for me would be a monad transformer that provides a supply of Refs, where data Ref v = Ref Int v, with the monad ensuring that Refs are given out with unique Int IDs, and Eq Ref etc. only looking at the Int (and now honesty is guaranteed by the uniqueness of the Ints). I would accept pointers to such a transformer in the wild as well.

user11228628
  • 1,526
  • 1
  • 6
  • 17
  • 2
    If you want to do lookups that cannot possibly fail how is that different from just having the values you were going to look up? – Cubic May 16 '19 at 08:13
  • 1
    Am I missing something obvious in asking why lookup values that you've just inserted, when you must still have access to the value? – moonGoose May 16 '19 at 08:16
  • 1
    The problem here might be that typically one constructs a map to lookup keys for which it is unknown whether these exist. Sure you can lookup `one` and `two`, but typically later on you use it with a stream of `Int`s, and thus the dictionary might fail. – Willem Van Onsem May 16 '19 at 08:43
  • How is holding on to `Receipt`s any different from holding on to the actual values? If you are somehow going to keep track of which `Receipt`s you will look up, that's equivalent to keeping track of what the looked-up value would be anyway (as @Cubic observed), so the lookup won't actually be a lookup in any meaningful sense of the word. I think a straightforward way to remove the unease of using `!` would be to use `ExceptT YourErrorType` with a `StateT (Map k v)` and throw and handle exceptions as you see fit. – aplainzetakind May 16 '19 at 10:10
  • 3
    @Cubic It is different because some other chunk of code which I don't know about could *also* insert that key later. Now when I go to look up using my receipt, I get their value, not the one I inserted. – Daniel Wagner May 16 '19 at 14:35
  • Then a simpler solution would be a `lookupWithDefault` with default being the inserted value you kept track of, though admittedly then you might use more memory than strictly necessary. – moonGoose May 16 '19 at 16:43
  • 1
    One possible motivation for this technique is avoiding ugly `Map.!` when traversing a graph defined as a map of adjacency lists (`Map nodeid [nodeid]`). – danidiaz May 16 '19 at 17:05
  • I've edited the question to clarify why a `Receipt` is better than holding the actual values. There is definitely some wiggle room in this solution space though, so thanks for the suggestions! – user11228628 May 16 '19 at 17:27
  • I don't think there's any existing implementation that *really* matches your use case. It just seems rather niche. There's nothing particularly wrong with your code, is there? All I can suggest is `StateT (Map k (STRef s v)) (ST s)` as another monad that supports similar operations. You don't have to do a `Map` lookup if you already have the `STRef` for a key, but you lose the ability to fork execution with e.g. `ListT`, because `ST` doesn't generate a well-behaved transformer. – HTNW May 17 '19 at 03:05

1 Answers1

3

Your solution resembles the technique used by justified-containers to guarantee that keys are present in a map. But there are some differences:

An expanded description of the technique used by justified-containers can be found in the functional pearl "Ghosts of departed proofs".

danidiaz
  • 26,936
  • 4
  • 45
  • 95
  • I like the design; I'm a little concerned about the usability of the continuations (relative to a monad transformer approach). If you have experience with this package, can you speak to how much of a burden it is to work with? – user11228628 May 16 '19 at 17:33
  • Actually, yeah, the recycling of the keys and the changing of the map's phantom type does present some difficulties for my use case. I'm repeatedly traversing a data structure until everything settles out, and inserting and looking up keys in the process, and I'm not sure how I'd do that loop if the types of the map and the keys keep changing. Maybe it's possible though. – user11228628 May 16 '19 at 17:46
  • @user11228628 I haven't used the library myself. – danidiaz May 16 '19 at 17:53