diff --git a/shardwidth/helper.go b/shardwidth/helper.go new file mode 100644 index 000000000..c04e2e72b --- /dev/null +++ b/shardwidth/helper.go @@ -0,0 +1,70 @@ +// Copyright 2021 Molecula Corp. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package shardwidth + +import ( + "math/bits" +) + +// FindNextShard returns the index of the first item which is not in the +// same shard as i. The index it returns may be equal to the length of the +// haystack, indicatincg that the rest of the list is in the same shard. +func FindNextShard(i int, haystack []uint64) int { + // compute the last thing that's in the same shard as haystack[i]. + if i >= len(haystack) { + return i + } + // current shard: + shard := (haystack[i] >> Exponent) + // last value in shard: + shardEnd := ((shard + 1) << Exponent) - 1 + j := i + // We want to do a binary search of the haystack. For any length of + // haystack, its topmost bit gives us a reasonable halfway point; it may + // not actually be halfway, but the number of steps it'll take to search + // it will be the same as if it were. sort.Search has interface overhead + // and makes us sad. + for incr := 1 << (bits.Len64(uint64(len(haystack) - i))); incr > 0; incr >>= 1 { + if j+incr < len(haystack) { + if haystack[j+incr] <= shardEnd { + j += incr + } + } + } + // we've found the last item that is in the same shard as i, so... + return j + 1 +} + +// FindShards finds the shards in a given haystack +func FindShards(haystack []uint64) (shards []uint64, endIndexes []int) { + if len(haystack) == 0 { + return nil, nil + } + index := 0 + // the steady state of this loop is that shards contains the current + // shard, but not its ending index; each time we find a new ending + // index, we record that index as the end for the current shard, and + // the new shard, until we reach the end and append len(haystack) + // as the last index. + shards = []uint64{haystack[index] >> Exponent} + index = FindNextShard(index, haystack) + for index < len(haystack) { + shards = append(shards, haystack[index]>>Exponent) + endIndexes = append(endIndexes, index) + index = FindNextShard(index, haystack) + } + endIndexes = append(endIndexes, index) + return shards, endIndexes +} diff --git a/shardwidth/helper_test.go b/shardwidth/helper_test.go new file mode 100644 index 000000000..cd1806cd8 --- /dev/null +++ b/shardwidth/helper_test.go @@ -0,0 +1,103 @@ +// Copyright 2021 Molecula Corp. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package shardwidth_test + +import ( + "math/rand" + "testing" + + "github.com/molecula/featurebase/v2/shardwidth" +) + +type nextShardTestCase struct { + name string + haystack [][2]uint64 // stored as shard, offset pairs + shardIndexes []int +} + +var nextShardTestCases = []nextShardTestCase{ + { + name: "all-in-one", + haystack: [][2]uint64{ + {0, 0}, + {0, 1}, + }, + shardIndexes: []int{2}, + }, + { + name: "split", + haystack: [][2]uint64{ + {0, 0}, + {1, 1}, + }, + shardIndexes: []int{1, 2}, + }, + { + name: "two-and-one", + haystack: [][2]uint64{ + {0, 0}, + {0, 1}, + {1, 1}, + }, + shardIndexes: []int{2, 3}, + }, +} + +func TestFindShards(t *testing.T) { + for _, c := range nextShardTestCases { + haystack := make([]uint64, len(c.haystack)) + for i, h := range c.haystack { + haystack[i] = (h[0] << shardwidth.Exponent) + h[1] + } + _, indexes := shardwidth.FindShards(haystack) + if len(indexes) != len(c.shardIndexes) { + t.Fatalf("%s: expected %d, got %d", c.name, c.shardIndexes, indexes) + } + for i, expected := range c.shardIndexes { + if indexes[i] != expected { + t.Fatalf("%s: expected index %d to be %d, got %d", c.name, i, expected, indexes[i]) + } + } + } + // fake up some more test cases + for i := 0; i < 100; i++ { + haystack := make([]uint64, 100) + shard := uint64(0) + bit := uint64(0) + shardIndexes := []int{} + for j := range haystack { + if rand.Intn(30) == 0 { + if j > 0 { + shardIndexes = append(shardIndexes, j) + } + shard++ + bit = 0 + } else { + bit += uint64(rand.Intn(30)) + } + haystack[j] = (shard << shardwidth.Exponent) + bit + } + shardIndexes = append(shardIndexes, len(haystack)) + _, indexes := shardwidth.FindShards(haystack) + if len(indexes) != len(shardIndexes) { + t.Fatalf("trial %d: expected %d, got %d", i, shardIndexes, indexes) + } + for idx, expected := range shardIndexes { + if indexes[idx] != expected { + t.Fatalf("trial %d: expected index %d to be %d, got %d", i, idx, expected, indexes[idx]) + } + } + } +}