Skip to content

Add NonEmpty containers, without imbalancing or massive code duplication #608

Open
@Ericson2314

Description

@Ericson2314

There's a number of packages on Hackage providing "non-empty" versions of containers, modeled after https://hackage.haskell.org/package/base/docs/Data-List-NonEmpty.html . There's a variety of use-cases, but perhaps my favorite is that they're needed to "curry" containers: for example, Map (k0, k1) v is isomorphic not to Map k0 (Map k1 v) but to Map k0 (NonEmptyMap k1 v). I like this use-case because it comes from the containers themselves.

Anyways, the problem with the existing implementations is that they tend to just prepend a minimal element to the container. This creates a less efficient, ill-balanced tree. But the alternative of copying the implementation creates a maintain burden. In https://github.com/mokus0/dependent-map/pull/31/files, I think I found an alternative which is the best of both worlds: mostly reused code and no degrade in balancing, namely making the types mutually recur.

For example, here are the types for DMap and NonEmptyDMap:

data DMap k f
  = Tip
  | Bin' {-# UNPACK #-} !(NonEmptyDMap k f)

pattern Bin s k v l r = Bin' (NonEmptyDMap s k v l r)

data NonEmptyDMap k f where
    NonEmptyDMap
         :: {- sz    -} !Int
         -> {- key   -} !(k v)
         -> {- value -} f v
         -> {- left  -} !(DMap k f)
         -> {- right -} !(DMap k f)
         -> NonEmptyDMap k f

DMap is like Maybe NonEmptyDMap but with the unboxed strict field.

Functions are then also written mutually recursively. (Take the exact weird worker wrapper style here with a grain of salt. This me trying to strike a balance between brevity and cargo culting the way the functions were written before.)

makeAdjustWithKey
  :: forall k f v
  .  GCompare k
  => (k v -> f v -> f v)
  -> k v
  -> ( DMap k f -> DMap k f
     , NonEmptyDMap k f -> NonEmptyDMap k f
     )
makeAdjustWithKey f k = (k `seq` go, k `seq` go')
  where
    go :: DMap k f -> DMap k f
    go Tip = Tip
    go (Bin' ne) = Bin' $! go' ne

    go' :: NonEmptyDMap k f -> NonEmptyDMap k f
    go' (NonEmptyDMap sx kx x l r) =
      case gcompare k kx of
        GLT -> NonEmptyDMap sx kx x (go l) r
        GGT -> NonEmptyDMap sx kx x l (go r)
        GEQ -> NonEmptyDMap sx kx (f kx x) l r
-- | /O(log n)/. Adjust a value at a specific key. When the key is not
-- a member of the map, the original map is returned.
adjustWithKey :: GCompare k => (k v -> f v -> f v) -> k v -> DMap k f -> DMap k f
adjustWithKey f k = fst $ makeAdjustWithKey f k
-- | /O(log n)/. Adjust a value at a specific key. When the key is not
-- a member of the map, the original map is returned.
adjustWithKey :: GCompare k => (k v -> f v -> f v) -> k v -> NonEmptyDMap k f -> NonEmptyDMap k f
adjustWithKey f k = snd $ makeAdjustWithKey f k

One interesting side benefit is that rotations and other internal operations can be less partial. Take rotateL, for example:

rotateL :: k v -> f v -> DMap k f -> NonEmptyDMap k f -> NonEmptyDMap k f
rotateL k x l r@(NonEmptyDMap _ _ _ ly ry)
  | sizeE ly < ratio*sizeE ry = singleL k x l r
  | otherwise                 = doubleL k x l r

The non-empty argument proves that at least one left rotation is possible. (The double rotate is still partial.)

As my parenthetical above hints, I haven't yet investigated performance properly; rather my goal was to do just enough work to demonstrate that something like this might be possible. If this sort of refactor looks viable to you, maintainers, I'd be happy to take a stab at doing it for all of containers, along with hunting down any performance issues that arise.

Metadata

Metadata

Assignees

No one assigned

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions