From f135f68aa002ac2b321bf7a4d1f3850e1040a8c0 Mon Sep 17 00:00:00 2001
From: hczphn <hczphn@polyhedra.network>
Date: Tue, 10 Dec 2024 19:16:00 -0600
Subject: [PATCH] add logup std

---
 circuit-std-ecgo/hint.go       |  47 ++++++
 circuit-std-ecgo/logup.go      | 281 +++++++++++++++++++++++++++++++++
 circuit-std-ecgo/logup_test.go | 126 +++++++++++++++
 circuit-std-ecgo/utils.go      | 129 +++++++++++++++
 4 files changed, 583 insertions(+)
 create mode 100644 circuit-std-ecgo/hint.go
 create mode 100644 circuit-std-ecgo/logup.go
 create mode 100644 circuit-std-ecgo/logup_test.go
 create mode 100644 circuit-std-ecgo/utils.go

diff --git a/circuit-std-ecgo/hint.go b/circuit-std-ecgo/hint.go
new file mode 100644
index 0000000..60f762c
--- /dev/null
+++ b/circuit-std-ecgo/hint.go
@@ -0,0 +1,47 @@
+package logup
+
+import (
+	"math/big"
+)
+
+func rangeProofHint(q *big.Int, inputs []*big.Int, outputs []*big.Int) error {
+	n := inputs[0].Int64()
+	a := new(big.Int).Set(inputs[1])
+
+	for i := int64(0); i < n/int64(LookupTableBits); i++ {
+		a, outputs[i] = new(big.Int).DivMod(a, big.NewInt(int64(1<<LookupTableBits)), new(big.Int))
+	}
+	return nil
+}
+
+// TODO: Do we need bits check for the count?
+func QueryCountHintFn(field *big.Int, inputs []*big.Int, outputs []*big.Int) error {
+	for i := 0; i < len(outputs); i++ {
+		outputs[i] = big.NewInt(0)
+	}
+
+	for i := 0; i < len(inputs); i++ {
+		query_id := inputs[i].Int64()
+		outputs[query_id].Add(outputs[query_id], big.NewInt(1))
+	}
+	return nil
+}
+
+func QueryCountBaseKeysHintFn(field *big.Int, inputs []*big.Int, outputs []*big.Int) error {
+	for i := 0; i < len(outputs); i++ {
+		outputs[i] = big.NewInt(0)
+	}
+	tableSize := inputs[0].Int64()
+	table := inputs[1 : tableSize+1]
+	queryKeys := inputs[tableSize+1:]
+	for i := 0; i < len(queryKeys); i++ {
+		queryKey := queryKeys[i].Int64()
+		//find the location of the query key in the table
+		for j := 0; j < len(table); j++ {
+			if table[j].Int64() == queryKey {
+				outputs[j].Add(outputs[j], big.NewInt(1))
+			}
+		}
+	}
+	return nil
+}
diff --git a/circuit-std-ecgo/logup.go b/circuit-std-ecgo/logup.go
new file mode 100644
index 0000000..938b413
--- /dev/null
+++ b/circuit-std-ecgo/logup.go
@@ -0,0 +1,281 @@
+package logup
+
+import (
+	"math"
+	"math/big"
+	"math/rand"
+
+	"github.com/consensys/gnark/constraint/solver"
+	"github.com/consensys/gnark/frontend"
+
+	"github.com/PolyhedraZK/ExpanderCompilerCollection/ecgo"
+)
+
+var (
+	Table           [][]frontend.Variable
+	QueryID         []frontend.Variable
+	QueryResult     [][]frontend.Variable
+	LookupTableBits int
+)
+
+func init() {
+	Table = make([][]frontend.Variable, 0)
+	QueryID = make([]frontend.Variable, 0)
+	QueryResult = make([][]frontend.Variable, 0)
+	Hints := []solver.Hint{rangeProofHint, QueryCountHintFn, QueryCountBaseKeysHintFn}
+	solver.RegisterHint(Hints...)
+}
+func Reset() {
+	LookupTableBits = 0
+	Table = make([][]frontend.Variable, 0)
+	QueryID = make([]frontend.Variable, 0)
+	QueryResult = make([][]frontend.Variable, 0)
+}
+
+// this interface write to a single (default) rangeProof table
+func NewRangeProof(bits int) {
+	LookupTableBits = bits
+	rangeLogupTableSize := 1 << bits
+	for i := 0; i < rangeLogupTableSize; i++ {
+		Table = append(Table, []frontend.Variable{i})
+	}
+}
+
+// this interface write to a single (default) table, the table size must be a power of 2
+func NewTable(key []frontend.Variable, value [][]frontend.Variable) {
+
+	if len(key) != len(value) {
+		panic("key and value should have the same length")
+	}
+	if !IsPowerOf2(len(key)) {
+		panic("the table size must be a power of 2")
+	}
+	if len(Table) != 0 {
+		panic("table should be empty")
+	}
+	//append key as the first column, and value as the rest columns
+	for i := 0; i < len(key); i++ {
+		entry := append([]frontend.Variable{key[i]}, value[i]...)
+		Table = append(Table, entry)
+	}
+}
+
+// this interface write a query to the default table, Table
+// a query is a pair of key and values
+func Query(key frontend.Variable, value []frontend.Variable) {
+	if len(value) != len(Table[0])-1 {
+		panic("value length should be equal to the table column size")
+	}
+	QueryID = append(QueryID, key)
+	//combine key and value as a query result
+	queryResult := append([]frontend.Variable{key}, value...)
+	QueryResult = append(QueryResult, queryResult)
+}
+
+func BatchQuery(keys []frontend.Variable, values [][]frontend.Variable) {
+	if len(keys) != len(values) || len(keys) == 0 {
+		panic("keys and values should have the same length and should not be empty")
+	}
+	if len(values[0]) != len(Table[0])-1 {
+		panic("value length should be equal to the table column size")
+	}
+	for i := 0; i < len(keys); i++ {
+		QueryID = append(QueryID, keys[i])
+		//combine key and value as a query result
+		queryResult := append([]frontend.Variable{keys[i]}, values[i]...)
+		QueryResult = append(QueryResult, queryResult)
+	}
+}
+
+// this interface write a query to the default table, Table
+// For a range query that checks if a key is in a table while ignoring the value (RangeProof)
+func QueryRange(key frontend.Variable) {
+	if len(Table[0]) != 1 {
+		panic("table should have only one column")
+	}
+	QueryID = append(QueryID, key)
+	QueryResult = append(QueryResult, []frontend.Variable{key})
+}
+
+func RangeProof(api frontend.API, a frontend.Variable, n int) {
+	//add a shift value
+	if n%LookupTableBits != 0 {
+		rem := n % LookupTableBits
+		shift := LookupTableBits - rem
+		constant := (1 << shift) - 1
+		mulFactor := big.NewInt(1)
+		mulFactor.Lsh(mulFactor, uint(n))
+		a = api.Add(a, api.Mul(constant, mulFactor))
+		n = n + shift
+	}
+	hintInput := make([]frontend.Variable, 2)
+	hintInput[0] = n
+	hintInput[1] = a
+	witnesses, err := api.Compiler().NewHint(rangeProofHint, n/LookupTableBits, hintInput...)
+	sum := witnesses[0]
+	for i := 1; i < len(witnesses); i++ {
+		constant := big.NewInt(1)
+		constant.Lsh(constant, uint(LookupTableBits*i))
+		sum = api.Add(sum, api.Mul(witnesses[i], constant))
+	}
+	api.AssertIsEqual(sum, a)
+	if err != nil {
+		panic(err)
+	}
+	for i := 0; i < n/LookupTableBits; i++ {
+		QueryRange(witnesses[i])
+	}
+}
+
+func FinalCheck(api frontend.Variable, column_combine_option ColumnCombineOptions) {
+	if len(Table) == 0 || len(QueryID) == 0 {
+		panic("empty table or empty query")
+	} // Should we allow this?
+	//if len(QueryID) != a power of 2, padding with query0
+	if !IsPowerOf2(len(QueryID)) {
+		nextPower2 := 1 << uint(math.Ceil(math.Log2(float64(len(QueryID)))))
+		for i := len(QueryID); i < nextPower2; i++ {
+			QueryID = append(QueryID, QueryID[0])
+			QueryResult = append(QueryResult, QueryResult[0])
+		}
+	}
+	ecgoApi := api.(ecgo.API)
+	// The challenge used to complete polynomial identity check
+	alpha := ecgoApi.GetRandomValue()
+
+	column_combine_randomness := GetColumnRandomness(ecgoApi, uint(len(Table[0])), column_combine_option)
+
+	// Table Polynomial
+	table_single_column := CombineColumn(ecgoApi, Table, column_combine_randomness)
+	// inputs :=
+	inputs := []frontend.Variable{frontend.Variable(len(Table))}
+	for i := 0; i < len(Table); i++ {
+		inputs = append(inputs, Table[i][0])
+	}
+	inputs = append(inputs, QueryID...)
+	query_count, _ := ecgoApi.NewHint(
+		QueryCountBaseKeysHintFn,
+		len(Table),
+		inputs...,
+	)
+
+	table_poly := make([]RationalNumber, len(table_single_column))
+	for i := 0; i < len(table_single_column); i++ {
+		table_poly[i] = RationalNumber{
+			Numerator:   query_count[i],
+			Denominator: ecgoApi.Sub(alpha, table_single_column[i]),
+		}
+	}
+	table_poly_at_alpha := SumRationalNumbers(ecgoApi, table_poly)
+	// Query Polynomial
+	query_single_column := CombineColumn(ecgoApi, QueryResult, column_combine_randomness)
+	query_poly := make([]RationalNumber, len(query_single_column))
+	for i := 0; i < len(query_single_column); i++ {
+		query_poly[i] = RationalNumber{
+			Numerator:   1,
+			Denominator: ecgoApi.Sub(alpha, query_single_column[i]),
+		}
+	}
+	query_poly_at_alpha := SumRationalNumbers(ecgoApi, query_poly)
+	ecgoApi.AssertIsEqual(
+		ecgoApi.Mul(table_poly_at_alpha.Numerator, query_poly_at_alpha.Denominator),
+		ecgoApi.Mul(query_poly_at_alpha.Numerator, table_poly_at_alpha.Denominator),
+	)
+}
+
+func NewRandomCircuit(
+	n_table_rows uint,
+	n_queries uint,
+	n_columns uint,
+	fill_values bool,
+) *LogUpCircuit {
+	c := &LogUpCircuit{}
+	c.Table = make([][]frontend.Variable, n_table_rows)
+	for i := 0; i < int(n_table_rows); i++ {
+		c.Table[i] = make([]frontend.Variable, n_columns)
+		if fill_values {
+			for j := 0; j < int(n_columns); j++ {
+				c.Table[i][j] = rand.Intn(math.MaxInt)
+			}
+		}
+	}
+
+	c.QueryID = make([]frontend.Variable, n_queries)
+	c.QueryResult = make([][]frontend.Variable, n_queries)
+
+	for i := 0; i < int(n_queries); i++ {
+		c.QueryResult[i] = make([]frontend.Variable, n_columns)
+		if fill_values {
+			query_id := rand.Intn(int(n_table_rows))
+			c.QueryID[i] = query_id
+			c.QueryResult[i] = c.Table[query_id]
+		}
+	}
+
+	return c
+}
+
+type ColumnCombineOptions int
+
+const (
+	Poly = iota
+	FullRandom
+)
+
+func (c *LogUpCircuit) Check(api ecgo.API, column_combine_option ColumnCombineOptions) error {
+	if len(c.Table) == 0 || len(c.QueryID) == 0 {
+		panic("empty table or empty query")
+	} // Should we allow this?
+
+	// The challenge used to complete polynomial identity check
+	alpha := api.GetRandomValue()
+
+	column_combine_randomness := GetColumnRandomness(api, uint(len(c.Table[0])), column_combine_option)
+
+	// Table Polynomial
+	table_single_column := CombineColumn(api, c.Table, column_combine_randomness)
+	query_count, _ := api.NewHint(
+		QueryCountHintFn,
+		len(c.Table),
+		c.QueryID...,
+	)
+
+	table_poly := make([]RationalNumber, len(table_single_column))
+	for i := 0; i < len(table_single_column); i++ {
+		table_poly[i] = RationalNumber{
+			Numerator:   query_count[i],
+			Denominator: api.Sub(alpha, table_single_column[i]),
+		}
+	}
+	table_poly_at_alpha := SumRationalNumbers(api, table_poly)
+
+	// Query Polynomial
+	query_single_column := CombineColumn(api, c.QueryResult, column_combine_randomness)
+	query_poly := make([]RationalNumber, len(query_single_column))
+	for i := 0; i < len(query_single_column); i++ {
+		query_poly[i] = RationalNumber{
+			Numerator:   1,
+			Denominator: api.Sub(alpha, query_single_column[i]),
+		}
+	}
+	query_poly_at_alpha := SumRationalNumbers(api, query_poly)
+
+	api.AssertIsEqual(
+		api.Mul(table_poly_at_alpha.Numerator, query_poly_at_alpha.Denominator),
+		api.Mul(query_poly_at_alpha.Numerator, table_poly_at_alpha.Denominator),
+	)
+	return nil
+}
+
+const ColumnCombineOption ColumnCombineOptions = FullRandom
+
+type LogUpCircuit struct {
+	Table       [][]frontend.Variable
+	QueryID     []frontend.Variable
+	QueryResult [][]frontend.Variable
+}
+
+// Define declares the circuit's constraints
+func (c *LogUpCircuit) Define(api frontend.API) error {
+	return c.Check(api.(ecgo.API), ColumnCombineOption)
+}
diff --git a/circuit-std-ecgo/logup_test.go b/circuit-std-ecgo/logup_test.go
new file mode 100644
index 0000000..c34dd3a
--- /dev/null
+++ b/circuit-std-ecgo/logup_test.go
@@ -0,0 +1,126 @@
+package logup
+
+import (
+	"os"
+	"testing"
+
+	"github.com/consensys/gnark-crypto/ecc"
+	"github.com/consensys/gnark/frontend"
+
+	"github.com/PolyhedraZK/ExpanderCompilerCollection/ecgo"
+	"github.com/PolyhedraZK/ExpanderCompilerCollection/ecgo/test"
+)
+
+func TestLogupCircuit(t *testing.T) {
+	N_TABLE_ROWS := uint(1024 * 4)
+	N_QUERIES := uint(1024 * 4)
+	COLUMN_SIZE := uint(2)
+
+	circuit, err := ecgo.Compile(ecc.BN254.ScalarField(), NewRandomCircuit(N_TABLE_ROWS, N_QUERIES, COLUMN_SIZE, false))
+	if err != nil {
+		panic(err.Error())
+	}
+
+	c := circuit.GetLayeredCircuit()
+	os.WriteFile("circuit.txt", c.Serialize(), 0o644)
+
+	assignment := NewRandomCircuit(N_TABLE_ROWS, N_QUERIES, COLUMN_SIZE, true)
+	inputSolver := circuit.GetInputSolver()
+	witness, err := inputSolver.SolveInput(assignment, 0)
+	if err != nil {
+		panic(err.Error())
+	}
+
+	if !test.CheckCircuit(c, witness) {
+		panic("Circuit not satisfied")
+	}
+
+	// os.WriteFile("inputsolver.txt", inputSolver.Serialize(), 0o644)
+	os.WriteFile("witness.txt", witness.Serialize(), 0o644)
+}
+
+type LogUpRuntimeCircuit struct {
+	Test [1]frontend.Variable
+}
+
+func (circuit *LogUpRuntimeCircuit) Define(api frontend.API) error {
+	// range proof
+	Reset()
+	keys := make([]frontend.Variable, 64)
+	values := make([][]frontend.Variable, 64)
+	for i := 0; i < 64; i++ {
+		keys[i] = frontend.Variable(i)
+		values[i] = []frontend.Variable{i + 1888}
+	}
+	NewTable(keys, values)
+	for i := 0; i < 64; i++ {
+		newI := (i + 23) % 64
+		Query(frontend.Variable(newI), []frontend.Variable{newI + 1888})
+	}
+	FinalCheck(api, ColumnCombineOption)
+	api.AssertIsEqual(frontend.Variable(1), circuit.Test[0])
+	return nil
+}
+func TestLogupRuntimeCircuit(t *testing.T) {
+	circuit, err := ecgo.Compile(ecc.BN254.ScalarField(), &LogUpRuntimeCircuit{})
+	if err != nil {
+		panic(err.Error())
+	}
+
+	c := circuit.GetLayeredCircuit()
+	os.WriteFile("circuit.txt", c.Serialize(), 0o644)
+
+	assignment := &LogUpRuntimeCircuit{Test: [1]frontend.Variable{1}}
+	inputSolver := circuit.GetInputSolver()
+	witness, err := inputSolver.SolveInput(assignment, 0)
+	if err != nil {
+		panic(err.Error())
+	}
+
+	if !test.CheckCircuit(c, witness) {
+		panic("Circuit not satisfied")
+	}
+
+	// os.WriteFile("inputsolver.txt", inputSolver.Serialize(), 0o644)
+	os.WriteFile("witness.txt", witness.Serialize(), 0o644)
+}
+
+type LogUpRangeProofCircuit struct {
+	Test [1]frontend.Variable
+}
+
+func (circuit *LogUpRangeProofCircuit) Define(api frontend.API) error {
+	// range proof
+	Reset()
+	NewRangeProof(8)
+	for i := 1; i < 12; i++ {
+		for j := (1 << (i - 1)); j < (1 << i); j++ {
+			RangeProof(api, frontend.Variable(j), i)
+		}
+	}
+	FinalCheck(api, ColumnCombineOption)
+	return nil
+}
+func TestLogupRangeProofCircuit(t *testing.T) {
+	circuit, err := ecgo.Compile(ecc.BN254.ScalarField(), &LogUpRangeProofCircuit{})
+	if err != nil {
+		panic(err.Error())
+	}
+
+	c := circuit.GetLayeredCircuit()
+	os.WriteFile("circuit.txt", c.Serialize(), 0o644)
+
+	assignment := &LogUpRangeProofCircuit{Test: [1]frontend.Variable{1}}
+	inputSolver := circuit.GetInputSolver()
+	witness, err := inputSolver.SolveInput(assignment, 0)
+	if err != nil {
+		panic(err.Error())
+	}
+
+	if !test.CheckCircuit(c, witness) {
+		panic("Circuit not satisfied")
+	}
+
+	// os.WriteFile("inputsolver.txt", inputSolver.Serialize(), 0o644)
+	os.WriteFile("witness.txt", witness.Serialize(), 0o644)
+}
diff --git a/circuit-std-ecgo/utils.go b/circuit-std-ecgo/utils.go
new file mode 100644
index 0000000..43ad75a
--- /dev/null
+++ b/circuit-std-ecgo/utils.go
@@ -0,0 +1,129 @@
+package logup
+
+import (
+	"fmt"
+
+	"github.com/consensys/gnark/frontend"
+
+	"github.com/PolyhedraZK/ExpanderCompilerCollection/ecgo"
+)
+
+type RationalNumber struct {
+	Numerator   frontend.Variable
+	Denominator frontend.Variable
+}
+
+func (r *RationalNumber) Add(api frontend.API, other *RationalNumber) RationalNumber {
+	return RationalNumber{
+		Numerator:   api.Add(api.Mul(r.Numerator, other.Denominator), api.Mul(other.Numerator, r.Denominator)),
+		Denominator: api.Mul(r.Denominator, other.Denominator),
+	}
+}
+
+// 0 is considered a power of 2 in this case
+func IsPowerOf2(n int) bool {
+	return n&(n-1) == 0
+}
+
+// Construct a binary summation tree to sum all the values
+func SumRationalNumbers(api frontend.API, rs []RationalNumber) RationalNumber {
+	n := len(rs)
+	if n == 0 {
+		return RationalNumber{Numerator: 0, Denominator: 1}
+	}
+
+	if !IsPowerOf2(n) {
+		fmt.Println(n)
+		panic("The length of rs should be a power of 2")
+	}
+
+	cur := rs
+	next := make([]RationalNumber, 0)
+
+	for n > 1 {
+		n >>= 1
+		for i := 0; i < n; i++ {
+			next = append(next, cur[i*2].Add(api, &cur[i*2+1]))
+		}
+		cur = next
+		next = next[:0]
+	}
+
+	if len(cur) != 1 {
+		panic("Summation code may be wrong.")
+	}
+
+	return cur[0]
+}
+
+func SimpleMin(a uint, b uint) uint {
+	if a < b {
+		return a
+	} else {
+		return b
+	}
+}
+
+func GetColumnRandomness(api ecgo.API, n_columns uint, column_combine_options ColumnCombineOptions) []frontend.Variable {
+	var randomness = make([]frontend.Variable, n_columns)
+	if column_combine_options == Poly {
+		beta := api.GetRandomValue()
+		randomness[0] = 1
+		randomness[1] = beta
+
+		// Hopefully this will generate fewer layers than sequential pows
+		max_deg := uint(1)
+		for max_deg < n_columns {
+			for i := max_deg + 1; i <= SimpleMin(max_deg*2, n_columns-1); i++ {
+				randomness[i] = api.Mul(randomness[max_deg], randomness[i-max_deg])
+			}
+			max_deg *= 2
+		}
+
+		// Debug Code:
+		// for i := 1; i < n_columns; i++ {
+		// 	api.AssertIsEqual(randomness[i], api.Mul(randomness[i - 1], beta))
+		// }
+
+	} else if column_combine_options == FullRandom {
+		randomness[0] = 1
+		for i := 1; i < int(n_columns); i++ {
+			randomness[i] = api.GetRandomValue()
+		}
+	} else {
+		panic("Unknown poly combine options")
+	}
+	return randomness
+}
+
+func CombineColumn(api ecgo.API, vec_2d [][]frontend.Variable, randomness []frontend.Variable) []frontend.Variable {
+	n_rows := len(vec_2d)
+	if n_rows == 0 {
+		return make([]frontend.Variable, 0)
+	}
+
+	n_columns := len(vec_2d[0])
+
+	// Do not introduce any randomness
+	if n_columns == 1 {
+		vec_combined := make([]frontend.Variable, n_rows)
+		for i := 0; i < n_rows; i++ {
+			vec_combined[i] = vec_2d[i][0]
+		}
+		return vec_combined
+	}
+
+	if !IsPowerOf2(n_columns) {
+		panic("Consider support this")
+	}
+
+	vec_return := make([]frontend.Variable, 0)
+	for i := 0; i < n_rows; i++ {
+		var v_at_row_i frontend.Variable = 0
+		for j := 0; j < n_columns; j++ {
+			v_at_row_i = api.Add(v_at_row_i, api.Mul(randomness[j], vec_2d[i][j]))
+		}
+		vec_return = append(vec_return, v_at_row_i)
+	}
+	return vec_return
+}