Monad transformer for inserts and total lookups on a Map?

Question

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.

Answer 1

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

• justified-containers uses continuation-passing-style.

• inserting a new key in justified-containers requires "recycling" existing receipts to work in the updated map. It seems that you don't need to "recyle" receipts because you never have multiple versions of the map around at the same time.

• justified-containers supports key deletion. I'm not sure if your solution could be expanded to allow the same.

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