poly.go

package main

type Poly struct{ C []FieldElement }

func NewPoly(coeffs []int64) Poly {
	c := make([]FieldElement, len(coeffs))
	for i, v := range coeffs {
		c[i] = NewFE(v)
	}
	return Poly{C: c}
}

func NewPolyFE(coeffs []FieldElement) Poly { return Poly{C: coeffs} }

func (p Poly) Degree() int {
	for i := len(p.C) - 1; i >= 0; i-- {
		if !p.C[i].IsZero() {
			return i
		}
	}
	return 0
}

func (p Poly) Trim() Poly {
	for i := len(p.C) - 1; i >= 0; i-- {
		if !p.C[i].IsZero() {
			return Poly{C: p.C[:i+1]}
		}
	}
	return Poly{}
}

func (p Poly) Add(o Poly) Poly {
	n := max(len(p.C), len(o.C))
	if n == 0 {
		return Poly{}
	}
	r := make([]FieldElement, n)
	for i := range r {
		r[i] = NewFE(0)
	}
	copy(r, p.C)
	for i, c := range o.C {
		r[i] = r[i].Add(c)
	}
	return NewPolyFE(r)
}

func (p Poly) MulScalar(s int64) Poly {
	se := NewFE(s)
	r := make([]FieldElement, len(p.C))
	for i, c := range p.C {
		r[i] = c.Mul(se)
	}
	return NewPolyFE(r)
}

func (p Poly) Mul(o Poly) Poly {
	if len(p.C) == 0 || len(o.C) == 0 {
		return Poly{}
	}
	r := make([]FieldElement, len(p.C)+len(o.C)-1)
	for i := range r {
		r[i] = NewFE(0)
	}
	for i, c1 := range p.C {
		for j, c2 := range o.C {
			r[i+j] = r[i+j].Add(c1.Mul(c2))
		}
	}
	return NewPolyFE(r)
}

func (p Poly) Div(d Poly) (Poly, Poly) {
	dd, dv := p.Degree(), d.Degree()
	if dd < dv {
		panic("invalid division")
	}
	dividend := make([]FieldElement, len(p.C))
	copy(dividend, p.C)
	qc := make([]FieldElement, dd-dv+1)
	for i := range qc {
		qc[i] = NewFE(0)
	}
	leadInv := d.C[dv].Inv()
	for i := len(qc) - 1; i >= 0; i-- {
		cd := dv + i
		if cd < len(dividend) && !dividend[cd].IsZero() {
			q := dividend[cd].Mul(leadInv)
			qc[i] = q
			for j := 0; j <= dv; j++ {
				if cd-dv+j < len(dividend) {
					dividend[cd-dv+j] = dividend[cd-dv+j].Sub(d.C[j].Mul(q))
				}
			}
		}
	}
	return NewPolyFE(qc), NewPolyFE(dividend).Trim()
}

func (p Poly) Eval(x FieldElement) FieldElement {
	r := NewFEWithField(0, x.F)
	for i, c := range p.C {
		r = r.Add(x.Pow(int64(i)).Mul(NewFEWithField(c.V, x.F)))
	}
	return r
}

func Interpolate(pts [][2]int64) Poly {
	n := len(pts)
	if n == 0 {
		return Poly{}
	}
	result := Poly{}
	for i := 0; i < n; i++ {
		xi, yi := pts[i][0], pts[i][1]
		basis := NewPoly([]int64{1})
		denom := NewFE(1)
		for j := 0; j < n; j++ {
			if i != j {
				basis = basis.Mul(NewPoly([]int64{-pts[j][0], 1}))
				denom = denom.Mul(NewFE(xi - pts[j][0]))
			}
		}
		result = result.Add(basis.MulScalar(NewFE(yi).Mul(denom.Inv()).V))
	}
	return result
}