Potential memory optimization for IntMap and IntSet

[meooow25]

This isn't particularly novel, so maybe it has been proposed before. Let me know if it has. I only found #340 in my search, which is a bigger idea.

Current definition

Today, Bin for IntMap and IntSet is represented as

containers/containers/src/Data/IntMap/Internal.hs

Lines 355 to 362 in 3c13e0b

	data IntMap a = Bin {-# UNPACK #-} !Prefix
	{-# UNPACK #-} !Mask
	!(IntMap a)
	!(IntMap a)
	-- Fields:
	-- prefix: The most significant bits shared by all keys in this Bin.
	-- mask: The switching bit to determine if a key should follow the left
	-- or right subtree of a 'Bin'.

Potential new definition

The prefix and the mask can be merged so that we save one word per Bin.

data IntMap a = Bin {-# UNPACK #-} !Int -- (current Prefix + current Mask)
                    !(IntMap a)
                    !(IntMap a)

The mask bit is always zero in the prefix, so this isn't throwing away any information. The lowest set bit of the new int is the current mask and the rest of it is the current prefix.

Current branching on Bin

Branching on Bin is currently done like this:

containers/containers/src/Data/IntMap/Internal.hs

Lines 813 to 817 in 3c13e0b

	insert :: Key -> a -> IntMap a -> IntMap a
	insert !k x t@(Bin p m l r)
	| nomatch k p m = link k (Tip k x) p t
	| zero k m = Bin p m (insert k x l) r
	| otherwise = Bin p m l (insert k x r)

containers/containers/src/Data/IntMap/Internal.hs

Lines 3527 to 3528 in 3c13e0b

	nomatch i p m
	= (mask i m) /= p

containers/containers/src/Data/IntMap/Internal.hs

Lines 3519 to 3520 in 3c13e0b

	zero i m
	= (natFromInt i) .&. (natFromInt m) == 0

New branching on Bin

insert :: Key -> a -> IntMap a -> IntMap a 
insert !k x t@(Bin pm l r) 
  | nomatch k pm = link k (Tip k x) p t 
  | left k pm    = Bin pm (insert k x l) r 
  | otherwise    = Bin pm l (insert k x r)
 
nomatch :: Int -> Int -> Bool
nomatch k pm = mask i pm /= pm .&. (pm-1)

left :: Int -> Int -> Bool
left k pm = int2word k < int2word pm 

Performance impact

I don't know yet, I need to make the change and benchmark. And it will involve changing every function, so it will be take a while.
Memory is certainly saved. nomatch gets a little more expensive, but left is cheaper than zero, so I'm hoping there is zero or positive overall effect.

What do you think? Is it worth checking out how this will fare? And is there any bad consequence of this representation, that I didn't think of?

[jwaldmann]

it will involve changing every function,

does it? transparent conversion between old and proposed new representation seems to me like the perfect use case for a pattern synonym - that GHC should be able to inline away.

I am skeptical of the overall effect on containers (certainly the original author(s) must have discussed this idea) but at least it could serve as a GHC test case.

[meooow25]

does it? transparent conversion between old and proposed new representation seems to me like the perfect use case for a pattern synonym - that GHC should be able to inline away.

Thanks, that should help incrementally test and benchmark it.

at least it could serve as a GHC test case.

Do you mean checking how it affects GHC tests?

---

Documenting some ideas: An alternative to the code above is making nomatch comparison based

i2w :: Int -> Word
i2w = fromIntegral

nomatch k pm = i2w k < i2w p || i2w (pm+m) < i2w k
  where
    p = pm .&. (pm-1)
    m = pm - p

In fact it can even be done today

nomatch k p m = i2w k < i2w p || i2w (p+m+m) < i2w k

I wonder if this is better than the current version, which, expanded, looks like

nomatch k p m = (k .&. (-m `xor` m)) /= p

[treeowl]

Definitely worth trying, if there really are enough bits! Give a pattern synonym a go; that should at least validate the semantics in a hurry. My experience with pattern synonyms is a little bit old. When I was messing around with them a bunch, I kept running into performance problems where they wouldn't inline well (IIRC, join point issues again). But now there are INLINE pragmas available for them, which might help.

[jwaldmann]

... performance problems with pattern synonyms

that's what I meant: the proposed transformation may serve as a test case for GHC's inliner.

[meooow25]

Cool, I'll test this out, maybe some time in the next couple of weeks.

[BurningWitness]

I'm hoping there is zero or positive overall effect.

I assume this will make the merges (and everything similar to them) worse. As mentioned nomatch becomes slower, and shorter has to retrieve each mask through n .&. (-n).

I wonder what the slowdown would be and whether bothering with it is even worth the meager 8 byte savings on each branch.

[meooow25]

I have tested out the changes on IntMap and I'm seeing modest improvements (5-15%) in lots of functions.
I'll prepare a PR. But for that, I would like to know if I should keep Bin as a pattern synonym for backwards compatibility (and rename the constructor to something else). I'm inclined to go with "no", because

1. There seems to be very few imports of Data.IntMap.Internal on Hackage. Seven, according to Serokell search: link. Not even all of them are using Bin.
2. Should someone use Bin, it is likely to get better performance for some operation. But a pattern synonym Bin that has to do some work would perform worse, and be counter to the author's intentions.

What do you think?

[treeowl]

I think keeping it as a pattern synonym would probably be a good thing, but we can decide later.