## Swapping

Posted on May 30, 2018
``````{-# 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           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.