correctly share buffer pointer across sampling instances

tuneinsight · Jun 11, 2024 · 0b1824f · 0b1824f
1 parent 3cc81c3
commit 0b1824f
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 9 deletions.
diff --git a/ring/sampler_gaussian.go b/ring/sampler_gaussian.go
@@ -5,6 +5,7 @@ import (
 	"math"
 	"math/big"
 
+	"github.com/tuneinsight/lattigo/v5/utils"
 	"github.com/tuneinsight/lattigo/v5/utils/bignum"
 	"github.com/tuneinsight/lattigo/v5/utils/sampling"
 )
@@ -18,7 +19,7 @@ type GaussianSampler struct {
 	baseSampler
 	xe            DiscreteGaussian
 	randomBufferN []byte
-	ptr           uint64
+	ptr           *uint64 // cross-instance buffer pointer
 	montgomery    bool
 }
 
@@ -29,7 +30,7 @@ func NewGaussianSampler(prng sampling.PRNG, baseRing *Ring, X DiscreteGaussian,
 	g = new(GaussianSampler)
 	g.prng = prng
 	g.randomBufferN = make([]byte, 1024)
-	g.ptr = 0
+	g.ptr = utils.Pointy[uint64](0)
 	g.baseRing = baseRing
 	g.xe = X
 	g.montgomery = montgomery
@@ -182,7 +183,7 @@ func (g *GaussianSampler) read(pol Poly, f func(a, b, c uint64) uint64) {
 // to use a secure PRNG instead of math/rand.
 func (g *GaussianSampler) normFloat64() (float64, uint64) {
 
-	ptr := g.ptr
+	ptr := *g.ptr
 	buff := g.randomBufferN
 	prng := g.prng
 	buffLen := uint64(len(buff))
@@ -224,7 +225,7 @@ func (g *GaussianSampler) normFloat64() (float64, uint64) {
 
 		// 1 (>99%)
 		if uint32(j) < kn[i] {
-			g.ptr = ptr
+			*g.ptr = ptr
 			return x, sign
 		}
 
@@ -242,13 +243,13 @@ func (g *GaussianSampler) normFloat64() (float64, uint64) {
 				}
 			}
 
-			g.ptr = ptr
+			*g.ptr = ptr
 			return x + rn, sign
 		}
 
 		// 3
 		if fn[i]+float32(randF64())*(fn[i-1]-fn[i]) < float32(math.Exp(-0.5*x*x)) {
-			g.ptr = ptr
+			*g.ptr = ptr
 			return x, sign
 		}
 	}

diff --git a/ring/sampler_uniform.go b/ring/sampler_uniform.go
@@ -11,7 +11,7 @@ import (
 type UniformSampler struct {
 	baseSampler
 	randomBufferN []byte
-	ptr           int
+	ptr           *int // cross-instance buffer pointer
 }
 
 // NewUniformSampler creates a new instance of UniformSampler from a PRNG and ring definition.
@@ -20,6 +20,7 @@ func NewUniformSampler(prng sampling.PRNG, baseRing *Ring) (u *UniformSampler) {
 	u.baseRing = baseRing
 	u.prng = prng
 	u.randomBufferN = make([]byte, utils.Max(1024, baseRing.N()))
+	u.ptr = utils.Pointy(0)
 	return
 }
 
@@ -55,7 +56,7 @@ func (u *UniformSampler) read(pol Poly, f func(a, b, c uint64) uint64) {
 	N := u.baseRing.N()
 
 	var ptr int
-	if ptr = u.ptr; ptr == 0 || ptr == N {
+	if ptr = *u.ptr; ptr == 0 || ptr == N {
 		if _, err := prng.Read(u.randomBufferN); err != nil {
 			// Sanity check, this error should not happen.
 			panic(err)
@@ -103,7 +104,7 @@ func (u *UniformSampler) read(pol Poly, f func(a, b, c uint64) uint64) {
 		}
 	}
 
-	u.ptr = ptr
+	*u.ptr = ptr
 }
 
 // ReadNew generates a new polynomial with coefficients following a uniform distribution over [0, Qi-1].