Swapping

Posted on May 30, 2018
Tags: 
{-# LANGUAGE RecursiveDo #-}

module Swap where

import qualified Data.Map.Strict as Map
import           Data.Map.Strict   (Map)

import           Data.IntMap          (IntMap)
import qualified Data.IntMap.Strict as IntMap
import qualified Data.IntMap.Lazy   as LazyIntMap

import           Control.Lens

import           Control.Arrow           ((&&&))
import           Control.Monad           ((>=>))
import           Control.Monad.Fix       (mfix)

import           Control.Monad.State     (StateT(..),execState,state)

import           Data.Maybe  (fromMaybe)
import           Data.Monoid (First(..))

Say you want to swap two items in a mapping structure—Data.Map.Strict, Data.HashMap, etc. The normal way uses far too many operations:

-- |
-- >>> swapAt4 1 2 (Map.fromList (zip [1..5] ['a'..]))
-- fromList [(1,'b'),(2,'a'),(3,'c'),(4,'d'),(5,'e')]
swapAt4 :: Ord a => a -> a -> Map a b -> Map a b
swapAt4 i j xs = case Map.lookup i xs of
  Nothing -> xs
  Just x -> case Map.lookup j xs of
    Nothing -> xs
    Just y -> Map.insert i y (Map.insert j x xs)

Two lookups, and two insertions. We can cut it down to three operations with insertLookupWithKey:

-- |
-- >>> swapAt3 1 2 (Map.fromList (zip [1..5] ['a'..]))
-- fromList [(1,'b'),(2,'a'),(3,'c'),(4,'d'),(5,'e')]
swapAt3 :: Ord a => a -> a -> Map a b -> Map a b
swapAt3 i j xs = case Map.lookup i xs of
  Nothing -> xs
  Just x -> case Map.insertLookupWithKey (const const) j x xs of
    (Nothing,_) -> xs
    (Just y,ys) -> Map.insert i y ys

Then, using laziness, we can write the above program circularly, reducing the number of lookups to 2:

swapAt2 :: Ord a => a -> a -> Map a b -> Map a b
swapAt2 i j xs = zs
  where
     (ival,ys) = Map.updateLookupWithKey (replace jval) i xs
     (jval,zs) = Map.updateLookupWithKey (replace ival) j ys
     replace x = const (Just . (`fromMaybe` x))

Unfortunately, Data.Map isn’t lazy enough for this: the above won’t terminate. Interestingly, Data.IntMap is lazy enough:

swapAt2Int :: Int -> Int -> IntMap a -> IntMap a
swapAt2Int i j xs = zs
  where
    (ival,ys) = LazyIntMap.updateLookupWithKey (replace jval) i xs
    (jval,zs) =     IntMap.updateLookupWithKey (replace ival) j ys
    replace x = const (Just . (`fromMaybe` x))

Notice how we have to use the lazy version of updateLookupWithKey. Again, though, this version has a problem: it won’t terminate when one of the keys is missing.

Thankfully, both of our problems can be solved by abstracting a little and using Ixed from lens:

-- |
-- >>> swapIx 1 2 "abc"
-- "acb"
swapIx :: Ixed a => Index a -> Index a -> a -> a
swapIx i j xs = zs
  where
    (First ival, ys) = ix i (replace jval) xs
    (First jval, zs) = ix j (replace ival) ys
    replace x = First . Just &&& (`fromMaybe` x)

Because ix is a traversal, it won’t do anything when there’s a missing key, which is what we want. Also, it adds extra laziness, as the caller of a traversal gets certain extra controls over the strictness of the traversal.

You may notice the stateful pattern above. However, translating it over as-is presents a problem: the circular bindings won’t work in vanilla do notation. For that, we need MonadFix and Recursive Do:

swapSt :: Ixed a => Index a -> Index a -> a -> a
swapSt i j = execState $ mdo
    ival <- replace i jval
    jval <- replace j ival
    pure ()
  where
    replace i (First x) =
        state (ix i (First . Just &&& (`fromMaybe` x)))

Finally, we can use mfix directly, and we’ll get the following clean-looking solution:

swap :: Ixed a => Index a -> Index a -> a -> a
swap i j = execState (mfix (replace i >=> replace j))
  where
    replace i (First x) =
        state (ix i (First . Just &&& (`fromMaybe` x)))

This works for most containers, even strict ones like Data.Map.Strict. It also works for Data.Vector. It does not work for Data.Vector.Unboxed, though.