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,99 +97,152 @@ proc intSetPut(t: var IntSet, key: int): PTrunk =
 
 proc contains*(s: IntSet, key: int): bool =
   ## returns true iff `key` is in `s`.
-  var t = intSetGet(s, `shr`(key, TrunkShift))
-  if t != nil:
-    var u = key and TrunkMask
-    result = (t.bits[`shr`(u, IntShift)] and `shl`(1, u and IntMask)) != 0
+  if s.elems <= s.a.len:
+    for i in 0..<s.elems:
+      if s.a[i] == key: return true
   else:
-    result = false
+    var t = intSetGet(s, `shr`(key, TrunkShift))
+    if t != nil:
+      var u = key and TrunkMask
+      result = (t.bits[`shr`(u, IntShift)] and `shl`(1, u and IntMask)) != 0
+    else:
+      result = false
 
-proc incl*(s: var IntSet, key: int) =
-  ## includes an element `key` in `s`.
+proc bitincl(s: var IntSet, key: int) {.inline.} =
   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`.
-  var t = intSetGet(s, `shr`(key, TrunkShift))
-  if t != nil:
-    var u = key and TrunkMask
-    t.bits[`shr`(u, IntShift)] = t.bits[`shr`(u, IntShift)] and
-        not `shl`(1, u and IntMask)
+  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
+      t.bits[`shr`(u, IntShift)] = t.bits[`shr`(u, IntShift)] and
+          not `shl`(1, u and IntMask)
 
 proc containsOrIncl*(s: var IntSet, key: int): bool =
   ## returns true if `s` contains `key`, otherwise `key` is included in `s`
   ## and false is returned.
-  var t = intSetGet(s, `shr`(key, TrunkShift))
-  if t != nil:
-    var u = key and TrunkMask
-    result = (t.bits[`shr`(u, IntShift)] and `shl`(1, u and IntMask)) != 0
-    if not result:
-      t.bits[`shr`(u, IntShift)] = t.bits[`shr`(u, IntShift)] or
-          `shl`(1, u and IntMask)
-  else:
+  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
+      result = (t.bits[`shr`(u, IntShift)] and `shl`(1, u and IntMask)) != 0
+      if not result:
+        t.bits[`shr`(u, IntShift)] = t.bits[`shr`(u, IntShift)] or
+            `shl`(1, u and IntMask)
+    else:
+      incl(s, key)
+      result = false
 
 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)
-  for i in 0..InitIntSetSize-1: result.data[i] = nil
-  result.max = InitIntSetSize-1
+  #setLen(result.data, InitIntSetSize)
+  #for i in 0..InitIntSetSize-1: result.data[i] = nil
+  #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`.
-  dest.counter = src.counter
-  dest.max = src.max
-  newSeq(dest.data, src.data.len)
+  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)
 
-  var it = src.head
-  while it != nil:
+    var it = src.head
+    while it != nil:
 
-    var h = it.key and dest.max
-    while dest.data[h] != nil: h = nextTry(h, dest.max)
-    assert(dest.data[h] == nil)
+      var h = it.key and dest.max
+      while dest.data[h] != nil: h = nextTry(h, dest.max)
+      assert(dest.data[h] == nil)
 
-    var n: PTrunk
-    new(n)
-    n.next = dest.head
-    n.key = it.key
-    n.bits = it.bits
-    dest.head = n
-    dest.data[h] = n
+      var n: PTrunk
+      new(n)
+      n.next = dest.head
+      n.key = it.key
+      n.bits = it.bits
+      dest.head = n
+      dest.data[h] = n
 
-    it = it.next
+      it = it.next
 
 iterator items*(s: IntSet): int {.inline.} =
   ## iterates over any included element of `s`.
-  var r = s.head
-  while r != nil:
-    var i = 0
-    while i <= high(r.bits):
-      var w = r.bits[i]
-      # taking a copy of r.bits[i] here is correct, because
-      # modifying operations are not allowed during traversation
-      var j = 0
-      while w != 0:         # test all remaining bits for zero
-        if (w and 1) != 0:  # the bit is set!
-          yield (r.key shl TrunkShift) or (i shl IntShift +% j)
-        inc(j)
-        w = w shr 1
-      inc(i)
-    r = r.next
+  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
+      while i <= high(r.bits):
+        var w = r.bits[i]
+        # taking a copy of r.bits[i] here is correct, because
+        # modifying operations are not allowed during traversation
+        var j = 0
+        while w != 0:         # test all remaining bits for zero
+          if (w and 1) != 0:  # the bit is set!
+            yield (r.key shl TrunkShift) or (i shl IntShift +% j)
+          inc(j)
+          w = w shr 1
+        inc(i)
+      r = r.next
 
 template dollarImpl(): untyped =
   result = "{"
@@ -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
+