Use Vose's Alias Method to sample the weighted distribution.

The weight generation code also was cleaned up (and now can support
generating distributions that look like what ScrambleSuit does as
a compile time change).

Per: http://www.keithschwarz.com/darts-dice-coins/

3 files changed, 225 insertions, 43 deletions

diff --git a/csrand/csrand.go b/csrand/csrand.go
index a3299aa..b059ed0 100644
--- a/csrand/csrand.go
+++ b/csrand/csrand.go
@@ -68,9 +68,9 @@ func (r csRandSource) Seed(seed int64) {
 	// No-op.
 }
 
-// Int63n returns, as a int64, a pseudo random number in [0, n).
-func Int63n(n int64) int64 {
-	return CsRand.Int63n(n)
+// Intn returns, as a int, a pseudo random number in [0, n).
+func Intn(n int) int {
+	return CsRand.Intn(n)
 }
 
 // Float64 returns, as a float64, a pesudo random number in [0.0,1.0).
diff --git a/weighted_dist.go b/weighted_dist.go
index 02fb26d..7c47cb8 100644
--- a/weighted_dist.go
+++ b/weighted_dist.go
@@ -28,6 +28,7 @@
 package obfs4
 
 import (
+	"container/list"
 	"fmt"
 	"math/rand"
 
@@ -36,27 +37,25 @@ import (
 )
 
 const (
-	minBuckets = 1
-	maxBuckets = 100
+	minValues = 1
+	maxValues = 100
 )
 
 // wDist is a weighted distribution.
 type wDist struct {
-	minValue    int
-	maxValue    int
-	values      []int
-	buckets     []int64
-	totalWeight int64
+	minValue int
+	maxValue int
+	values   []int
+	weights  []float64
 
-	rng *rand.Rand
+	alias []int
+	prob  []float64
 }
 
 // newWDist creates a weighted distribution of values ranging from min to max
 // based on a HashDrbg initialized with seed.
 func newWDist(seed *drbg.Seed, min, max int) (w *wDist) {
-	w = new(wDist)
-	w.minValue = min
-	w.maxValue = max
+	w = &wDist{minValue: min, maxValue: max}
 
 	if max <= min {
 		panic(fmt.Sprintf("wDist.Reset(): min >= max (%d, %d)", min, max))
@@ -67,46 +66,147 @@ func newWDist(seed *drbg.Seed, min, max int) (w *wDist) {
 	return
 }
 
-// sample generates a random value according to the distribution.
-func (w *wDist) sample() int {
-	retIdx := 0
-	var totalWeight int64
-	weight := csrand.Int63n(w.totalWeight)
-	for i, bucketWeight := range w.buckets {
-		totalWeight += bucketWeight
-		if weight <= totalWeight {
-			retIdx = i
-			break
+// genValues creates a slice containing a random number of random values
+// that when scaled by adding minValue will fall into [min, max].
+func (w *wDist) genValues(rng *rand.Rand) {
+	nValues := (w.maxValue + 1) - w.minValue
+	values := rng.Perm(nValues)
+	if nValues < minValues {
+		nValues = minValues
+	}
+	if nValues > maxValues {
+		nValues = maxValues
+	}
+	nValues = rng.Intn(nValues) + 1
+	w.values = values[:nValues]
+}
+
+// genBiasedWeights generates a non-uniform weight list, similar to the
+// ScrambleSuit prob_dist module.
+func (w *wDist) genBiasedWeights(rng *rand.Rand) {
+	w.weights = make([]float64, len(w.values))
+
+	culmProb := 0.0
+	for i := range w.values {
+		p := (1.0 - culmProb) * rng.Float64()
+		w.weights[i] = p
+		culmProb += p
+	}
+}
+
+// genUniformWeights generates a uniform weight list.
+func (w *wDist) genUniformWeights(rng *rand.Rand) {
+	w.weights = make([]float64, len(w.values))
+	for i := range w.weights {
+		w.weights[i] = rng.Float64()
+	}
+}
+
+// genTables calculates the alias and prob tables used for Vose's Alias method.
+// Algorithm taken from http://www.keithschwarz.com/darts-dice-coins/
+func (w *wDist) genTables() {
+	n := len(w.weights)
+	var sum float64
+	for _, weight := range w.weights {
+		sum += weight
+	}
+
+	// Create arrays $Alias$ and $Prob$, each of size $n$.
+	alias := make([]int, n)
+	prob := make([]float64, n)
+
+	// Create two worklists, $Small$ and $Large$.
+	small := list.New()
+	large := list.New()
+
+	scaled := make([]float64, n)
+	for i, weight := range w.weights {
+		// Multiply each probability by $n$.
+		p_i := weight * float64(n) / sum
+		scaled[i] = p_i
+
+		// For each scaled probability $p_i$:
+		if scaled[i] < 1.0 {
+			// If $p_i < 1$, add $i$ to $Small$.
+			small.PushBack(i)
+		} else {
+			// Otherwise ($p_i \ge 1$), add $i$ to $Large$.
+			large.PushBack(i)
 		}
 	}
 
-	return w.minValue + w.values[retIdx]
+	// While $Small$ and $Large$ are not empty: ($Large$ might be emptied first)
+	for small.Len() > 0 && large.Len() > 0 {
+		// Remove the first element from $Small$; call it $l$.
+		l := small.Remove(small.Front()).(int)
+		// Remove the first element from $Large$; call it $g$.
+		g := large.Remove(large.Front()).(int)
+
+		// Set $Prob[l] = p_l$.
+		prob[l] = scaled[l]
+		// Set $Alias[l] = g$.
+		alias[l] = g
+
+		// Set $p_g := (p_g + p_l) - 1$. (This is a more numerically stable option.)
+		scaled[g] = (scaled[g] + scaled[l]) - 1.0
+
+		if scaled[g] < 1.0 {
+			// If $p_g < 1$, add $g$ to $Small$.
+			small.PushBack(g)
+		} else {
+			// Otherwise ($p_g \ge 1$), add $g$ to $Large$.
+			large.PushBack(g)
+		}
+	}
+
+	// While $Large$ is not empty:
+	for large.Len() > 0 {
+		// Remove the first element from $Large$; call it $g$.
+		g := large.Remove(large.Front()).(int)
+		// Set $Prob[g] = 1$.
+		prob[g] = 1.0
+	}
+
+	// While $Small$ is not empty: This is only possible due to numerical instability.
+	for small.Len() > 0 {
+		// Remove the first element from $Small$; call it $l$.
+		l := small.Remove(small.Front()).(int)
+		// Set $Prob[l] = 1$.
+		prob[l] = 1.0
+	}
+
+	w.prob = prob
+	w.alias = alias
 }
 
 // reset generates a new distribution with the same min/max based on a new seed.
 func (w *wDist) reset(seed *drbg.Seed) {
 	// Initialize the deterministic random number generator.
 	drbg := drbg.NewHashDrbg(seed)
-	w.rng = rand.New(drbg)
+	rng := rand.New(drbg)
 
-	nBuckets := (w.maxValue + 1) - w.minValue
-	w.values = w.rng.Perm(nBuckets)
-	if nBuckets < minBuckets {
-		nBuckets = minBuckets
-	}
-	if nBuckets > maxBuckets {
-		nBuckets = maxBuckets
-	}
-	nBuckets = w.rng.Intn(nBuckets) + 1
-
-	w.totalWeight = 0
-	w.buckets = make([]int64, nBuckets)
-	for i, _ := range w.buckets {
-		prob := w.rng.Int63n(1000)
-		w.buckets[i] = prob
-		w.totalWeight += prob
+	w.genValues(rng)
+	//w.genBiasedWeights(rng)
+	w.genUniformWeights(rng)
+	w.genTables()
+}
+
+// sample generates a random value according to the distribution.
+func (w *wDist) sample() int {
+	var idx int
+
+	// Generate a fair die roll from an $n$-sided die; call the side $i$.
+	i := csrand.Intn(len(w.values))
+	// Flip a biased coin that comes up heads with probability $Prob[i]$.
+	if csrand.Float64() <= w.prob[i] {
+		// If the coin comes up "heads," return $i$.
+		idx = i
+	} else {
+		// Otherwise, return $Alias[i]$.
+		idx = w.alias[i]
 	}
-	w.buckets[len(w.buckets)-1] = w.totalWeight
+
+	return w.minValue + w.values[idx]
 }
 
 /* vim :set ts=4 sw=4 sts=4 noet : */
diff --git a/weighted_dist_test.go b/weighted_dist_test.go
new file mode 100644
index 0000000..14fecec
--- /dev/null
+++ b/weighted_dist_test.go
@@ -0,0 +1,82 @@
+/*
+ * Copyright (c) 2014, Yawning Angel <yawning at torproject dot org>
+ * All rights reserved.
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ *  * Redistributions of source code must retain the above copyright notice,
+ *    this list of conditions and the following disclaimer.
+ *
+ *  * Redistributions in binary form must reproduce the above copyright notice,
+ *    this list of conditions and the following disclaimer in the documentation
+ *    and/or other materials provided with the distribution.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+ * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+ * ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
+ * LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
+ * CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+ * SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+ * INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
+ * CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+ * ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+ * POSSIBILITY OF SUCH DAMAGE.
+ */
+
+package obfs4
+
+import (
+	"fmt"
+	"testing"
+
+	"github.com/yawning/obfs4/drbg"
+)
+
+const debug = false
+
+func TestWeightedDist(t *testing.T) {
+	seed, err := drbg.NewSeed()
+	if err != nil {
+		t.Fatal("failed to generate a DRBG seed:", err)
+	}
+
+	const nrTrials = 1000000
+
+	hist := make([]int, 1000)
+
+	w := newWDist(seed, 0, 999)
+	if debug {
+		// Dump a string representation of the probability table.
+		fmt.Println("Table:")
+		var sum float64
+		for _, weight := range w.weights {
+			sum += weight
+		}
+		for i, weight := range w.weights {
+			p := weight / sum
+			if p > 0.000001 { // Filter out tiny values.
+				fmt.Printf(" [%d]: %f\n", w.minValue+w.values[i], p)
+			}
+		}
+		fmt.Println()
+	}
+
+	for i := 0; i < nrTrials; i++ {
+		value := w.sample()
+		hist[value]++
+	}
+
+	if debug {
+		fmt.Println("Generated:")
+		for value, count := range hist {
+			if count != 0 {
+				p := float64(count) / float64(nrTrials)
+				fmt.Printf(" [%d]: %f (%d)\n", value, p, count)
+			}
+		}
+	}
+}
+
+/* vim :set ts=4 sw=4 sts=4 noet : */