optimized intsets to not allocate for the common cases

lib/pure/collections/intsets.nim

@@ -31,16 +31,18 @@ const
 
 type
   PTrunk = ref Trunk
-  Trunk {.final.} = object
+  Trunk = object
     next: PTrunk             # all nodes are connected with this pointer
     key: int                 # start address at bit 0
     bits: array[0..IntsPerTrunk - 1, BitScalar] # a bit vector
 
   TrunkSeq = seq[PTrunk]
   IntSet* = object ## an efficient set of 'int' implemented as a sparse bit set
+    elems: int # only valid for small numbers
     counter, max: int
     head: PTrunk
     data: TrunkSeq
+    a: array[0..33, int] # profiling shows that 34 elements are enough
 
 {.deprecated: [TIntSet: IntSet, TTrunk: Trunk, TTrunkSeq: TrunkSeq].}
 
@@ -95,6 +97,10 @@ proc intSetPut(t: var IntSet, key: int): PTrunk =
 
 proc contains*(s: IntSet, key: int): bool =
   ## returns true iff `key` is in `s`.
+  if s.elems <= s.a.len:
+    for i in 0..<s.elems:
+      if s.a[i] == key: return true
+  else:
     var t = intSetGet(s, `shr`(key, TrunkShift))
     if t != nil:
       var u = key and TrunkMask
@@ -102,15 +108,38 @@ proc contains*(s: IntSet, key: int): bool =
     else:
       result = false
 
-proc incl*(s: var IntSet, key: int) =
+proc bitincl(s: var IntSet, key: int) {.inline.} =
-  ## includes an element `key` in `s`.
   var t = intSetPut(s, `shr`(key, TrunkShift))
   var u = key and TrunkMask
   t.bits[`shr`(u, IntShift)] = t.bits[`shr`(u, IntShift)] or
       `shl`(1, u and IntMask)
 
+proc incl*(s: var IntSet, key: int) =
+  ## includes an element `key` in `s`.
+  if s.elems <= s.a.len:
+    for i in 0..<s.elems:
+      if s.a[i] == key: return
+    if s.elems < s.a.len:
+      s.a[s.elems] = key
+      inc s.elems
+      return
+    newSeq(s.data, InitIntSetSize)
+    s.max = InitIntSetSize-1
+    for i in 0..<s.elems:
+      bitincl(s, s.a[i])
+    s.elems = s.a.len + 1
+    # fall through:
+  bitincl(s, key)
+
 proc excl*(s: var IntSet, key: int) =
   ## excludes `key` from the set `s`.
+  if s.elems <= s.a.len:
+    for i in 0..<s.elems:
+      if s.a[i] == key:
+        s.a[i] = s.a[s.elems-1]
+        dec s.elems
+        return
+  else:
     var t = intSetGet(s, `shr`(key, TrunkShift))
     if t != nil:
       var u = key and TrunkMask
@@ -120,6 +149,13 @@ proc excl*(s: var IntSet, key: int) =
 proc containsOrIncl*(s: var IntSet, key: int): bool =
   ## returns true if `s` contains `key`, otherwise `key` is included in `s`
   ## and false is returned.
+  if s.elems <= s.a.len:
+    for i in 0..<s.elems:
+      if s.a[i] == key:
+        return true
+    incl(s, key)
+    result = false
+  else:
     var t = intSetGet(s, `shr`(key, TrunkShift))
     if t != nil:
       var u = key and TrunkMask
@@ -133,23 +169,38 @@ proc containsOrIncl*(s: var IntSet, key: int): bool =
 
 proc initIntSet*: IntSet =
   ## creates a new int set that is empty.
-  newSeq(result.data, InitIntSetSize)
+
-  result.max = InitIntSetSize-1
+  #newSeq(result.data, InitIntSetSize)
+  #result.max = InitIntSetSize-1
+  result.data = nil
+  result.max = 0
   result.counter = 0
   result.head = nil
+  result.elems = 0
 
 proc clear*(result: var IntSet) =
-  setLen(result.data, InitIntSetSize)
+  #setLen(result.data, InitIntSetSize)
-  for i in 0..InitIntSetSize-1: result.data[i] = nil
+  #for i in 0..InitIntSetSize-1: result.data[i] = nil
-  result.max = InitIntSetSize-1
+  #result.max = InitIntSetSize-1
+  result.data = nil
+  result.max = 0
   result.counter = 0
   result.head = nil
+  result.elems = 0
 
-proc isNil*(x: IntSet): bool {.inline.} = x.head.isNil
+proc isNil*(x: IntSet): bool {.inline.} = x.head.isNil and x.elems == 0
 
 proc assign*(dest: var IntSet, src: IntSet) =
   ## copies `src` to `dest`. `dest` does not need to be initialized by
   ## `initIntSet`.
+  if src.elems <= src.a.len:
+    dest.data = nil
+    dest.max = 0
+    dest.counter = src.counter
+    dest.head = nil
+    dest.elems = src.elems
+    dest.a = src.a
+  else:
     dest.counter = src.counter
     dest.max = src.max
     newSeq(dest.data, src.data.len)
@@ -173,6 +224,10 @@ proc assign*(dest: var IntSet, src: IntSet) =
 
 iterator items*(s: IntSet): int {.inline.} =
   ## iterates over any included element of `s`.
+  if s.elems <= s.a.len:
+    for i in 0..<s.elems:
+      yield s.a[i]
+  else:
     var r = s.head
     while r != nil:
       var i = 0
@@ -225,3 +280,9 @@ when isMainModule:
   ys.sort(cmp[int])
   assert ys == @[1, 2, 7, 1056]
 
+  var z: IntSet
+  for i in 0..1000:
+    incl z, i
+  for i in 0..1000:
+    assert i in z
+