schema

type AnySchema = Schema[Constraint]

type Assignment interface {
	// For the given module, determine any well-definedness bounds implied by
	// this assignment in  both the negative (left) or positive (right)
	// directions.  For example, consider an expression such as "(shift X -1)".
	// This is technically undefined for the first row of any trace and, by
	// association, any assignment evaluating this expression on that first row
	// is also undefined.
	Bounds(ModuleId) util.Bounds
	// ComputeColumns computes the values of columns defined by this assignment.
	// In order for this computation to makes sense, all columns on which this
	// assignment depends must exist (e.g. are either inputs or have been
	// computed already).  Computed columns do not exist in the original trace,
	// but are added during trace expansion to form the final trace.
	Compute(tr.Trace[bls12_377.Element], Schema[Constraint]) ([]tr.ArrayColumn, error)
	// Consistent applies a number of internal consistency checks.  Whilst not
	// strictly necessary, these can highlight otherwise hidden problems as an aid
	// to debugging.
	Consistent(Schema[Constraint]) []error
	// Identifier registers which are expanded by this assignment.  A register
	// is expanded when its length maybe changed.  For example, when going from
	// a trace which contains only rows of input/output values to a trace where
	// each function instance can occupy more than one row.  Then the I/O
	// columns are said to be "expanded".
	RegistersExpanded() []RegisterRef
	// Returns the set of columns that this assignment depends upon.  That can
	// include both input columns, as well as other computed columns.
	RegistersRead() []RegisterRef
	// Identifier registers assigned by this assignment.
	RegistersWritten() []RegisterRef
	// Lisp converts this schema element into a simple S-Expression, for example
	// so it can be printed.
	Lisp(AnySchema) sexp.SExp
}

type Constraint interface {
	// Accepts determines whether a given constraint accepts a given trace or
	// not.  If not, a failure is produced.  Otherwise, a bitset indicating
	// branch coverage is returned.
	Accepts(trace.Trace[bls12_377.Element], AnySchema) (bit.Set, Failure)
	// Determine the well-definedness bounds for this constraint in both the
	// negative (left) or positive (right) directions.  For example, consider an
	// expression such as "(shift X -1)".  This is technically undefined for the
	// first row of any trace and, by association, any constraint evaluating
	// this expression on that first row is also undefined (and hence must pass)
	Bounds(module uint) util.Bounds
	// Consistent applies a number of internal consistency checks.  Whilst not
	// strictly necessary, these can highlight otherwise hidden problems as an aid
	// to debugging.
	Consistent(Schema[Constraint]) []error
	// Contexts returns the evaluation contexts (i.e. enclosing module + length
	// multiplier) for this constraint.  Most constraints have only a single
	// evaluation context, though some (e.g. lookups) have more.  Note that all
	// constraints have at least one context (which we can call the "primary"
	// context).
	Contexts() []ModuleId
	// Name returns a unique name for a given constraint.  This is useful purely
	// for identifying constraints in reports, etc.
	Name() string
	// Lisp converts this schema element into a simple S-Expression, for example
	// so it can be printed.
	Lisp(AnySchema) sexp.SExp
	// Substitute any matchined labelled constants within this constraint
	Substitute(map[string]fr.Element)
}

type Expander[M any, C any] func(Schema[C], trace.Trace[bls12_377.Element]) trace.Trace[bls12_377.Element]

type Failure interface {
	// Provides a suitable error message
	Message() string
}

type FieldAgnostic[T any] interface {
	// Subdivide for a given bandwidth and maximum register width. This will
	// split all registers wider than the maximum permitted width into two or
	// more "limbs" (i.e. subregisters which do not exceeded the permitted
	// width).  For example, consider a register "r" of width u32. Subdividing
	// this register into registers of at most 8bits will result in four limbs:
	// r'0, r'1, r'2 and r'3 where (by convention) r'0 is the least significant.
	//
	// As part of the subdivision process, constraints may also need to be
	// divided when they exceed the maximum permitted bandwidth.  For example,
	// consider a simple constraint such as "x = y + 1" using 16bit registers
	// x,y.  Subdividing for a bandwidth of 10bits and a maximum register width
	// of 8bits means splitting each register into two limbs, and transforming
	// our constraint into:
	//
	// 256*x'1 + x'0 = 256*y'1 + y'0 + 1
	//
	// However, as it stands, this constraint exceeds our bandwidth requirement
	// since it requires at least 17bits of information to safely evaluate each
	// side.  Thus, the constraint itself must be subdivided into two parts:
	//
	// 256*c + x'0 = y'0 + 1  // lower
	//
	//         x'1 = y'1 + c  // upper
	//
	// Here, c is a 1bit register introduced as part of the transformation to
	// act as a "carry" between the two constraints.
	Subdivide(LimbsMap) T
}

type FieldAgnosticModule[M Module] interface {
	Module
	FieldAgnostic[M]
}

type FieldConfig struct {
	// Name suitable for identifying the config.  This is only really used for
	// improving error reporting, etc.
	Name string
	// Maximum field bandwidth available in the field.
	FieldBandWidth uint
	// Maximum register width to use with this field.
	RegisterWidth uint
}

type Limb = Register

type LimbId = RegisterId

type LimbsMap = ModuleMap[RegisterLimbsMap]

type MixedSchema[M1 Module, M2 Module] struct {
	// contains filtered or unexported fields
}

type Module interface {
	RegisterMap
	// Assignments returns an iterator over the assignments of this module.
	// These are the computations used to assign values to all computed columns
	// in this module.
	Assignments() iter.Iterator[Assignment]
	// Constraints provides access to those constraints associated with this
	// module.
	Constraints() iter.Iterator[Constraint]
	// Consistent applies a number of internal consistency checks.  Whilst not
	// strictly necessary, these can highlight otherwise hidden problems as an aid
	// to debugging.
	Consistent(Schema[Constraint]) []error
	// Identifies the length multiplier for this module.  For every trace, the
	// height of the corresponding module must be a multiple of this.  This is
	// used specifically to support interleaving constraints.
	LengthMultiplier() uint
	// AllowPadding determines the amount of initial padding a module expects.
	AllowPadding() bool
	// Module name
	Name() string
	// Returns the number of registers in this module.
	Width() uint
}

type ModuleId = uint

type ModuleMap[T RegisterMap] interface {
	// Field returns the underlying field configuration used for this mapping.
	// This includes the field bandwidth (i.e. number of bits available in
	// underlying field) and the maximum register width (i.e. width at which
	// registers are capped).
	Field() FieldConfig
	// Module returns register mapping information for the given module.
	Module(ModuleId) T
	// ModuleOf returns register mapping information for the given module.
	ModuleOf(string) T
}

type Register struct {
	// Kind of register (input / output)
	Kind RegisterType
	// Given name of this register.
	Name string
	// Width (in bits) of this register
	Width uint
}

type RegisterAllocator interface {
	RegisterLimbsMap
	// AllocateCarry a fresh register of the given width within the target module.
	// This is presumed to be a computed register, and automatically assigned a
	// unique name.
	AllocateCarry(prefix string, width uint) RegisterId
}

type RegisterId = trace.ColumnId

type RegisterLimbsMap interface {
	RegisterMap
	// Field returns the underlying field configuration used for this mapping.
	// This includes the field bandwidth (i.e. number of bits available in
	// underlying field) and the maximum register width (i.e. width at which
	// registers are capped).
	Field() FieldConfig
	// Limbs identifies the limbs into which a given register is divided.
	// Observe that limbs are ordered by their position in the original
	// register.  In particular, the first limb (i.e. at index 0) is always
	// least significant limb, and the last always most significant.
	LimbIds(RegisterId) []LimbId
	// Limbs returns information about a given limb (i.e. a register which
	// exists after the split).
	Limb(LimbId) Limb
	// Limbs returns all limbs in the mapping.
	Limbs() []Limb
	// LimbsMap returns a register map for the limbs themselves.  This is useful
	// where we need a register map over the limbs, rather than the original
	// registers.
	LimbsMap() RegisterMap
}

type RegisterMap interface {
	// Name returns the name given to the enclosing entity (i.e. module or
	// function).
	Name() string
	// HasRegister checks whether a register with the given name exists and, if
	// so, returns its register identifier.  Otherwise, it returns false.
	HasRegister(name string) (RegisterId, bool)
	// Access a given register in this module.
	Register(RegisterId) Register
	// Registers providers access to the underlying registers of this map.
	Registers() []Register
}

type RegisterRef = trace.ColumnRef

type RegisterType struct {
	// contains filtered or unexported fields
}

type Schema[C any] interface {
	// Assignments returns an iterator over the assignments of this schema.
	// That is, the set of computations used to determine values for all
	// computed columns.
	Assignments() iter.Iterator[Assignment]
	// Consistent applies a number of internal consistency checks.  Whilst not
	// strictly necessary, these can highlight otherwise hidden problems as an aid
	// to debugging.
	Consistent() []error
	// Constraints returns an iterator over all constraints defined in this
	// schema.  Observe that this does include assertions which, strictly
	// speaking, are not constraints in the true sense.  That is, they are never
	// compiled into vanishing polynomials but, instead, are used purely for
	// debugging.
	Constraints() iter.Iterator[C]
	// HasModule checks whether a module with the given name exists and, if so,
	// returns its module identifier.  Otherwise, it returns false.
	HasModule(name string) (ModuleId, bool)
	// Access a given module in this schema.
	Module(module uint) Module
	// Modules returns an iterator over the declared set of modules within this
	// schema.
	Modules() iter.Iterator[Module]
	// Access a given register in this schema.
	Register(RegisterRef) Register
	// Returns the number of modules in this schema.
	Width() uint
}

type Table[C Constraint] struct {
	// contains filtered or unexported fields
}

type UniformSchema[M Module] struct {
	// contains filtered or unexported fields
}