Feat/roman/tree builder (#525)

# Updates:

## Hashing

 - Added SpongeHasher class
 - Can be used to accept any hash function as an argument
 - Absorb and squeeze are now separated
- Memory management is now mostly done by SpongeHasher class, each hash
function only describes permutation kernels

## Tree builder

 - Tree builder is now hash-agnostic. 
 - Tree builder now supports 2D input (matrices)
- Tree builder can now use two different hash functions for layer 0 and
compression layers

## Poseidon1

 - Interface changed to classes
 - Now allows for any alpha
 - Now allows passing constants not in a single vector
 - Now allows for any domain tag
 - Constants are now released upon going out of scope
 - Rust wrappers changed to Poseidon struct
 
 ## Poseidon2
 
 - Interface changed to classes
 - Constants are now released upon going out of scope
 - Rust wrappers changed to Poseidon2 struct
 
## Keccak

 - Added Keccak class which inherits SpongeHasher
 - Now doesn't use gpu registers for storing states
 
 To do:
- [x] Update poseidon1 golang bindings
- [x] Update poseidon1 examples
- [x] Fix poseidon2 cuda test
- [x] Fix poseidon2 merkle tree builder test
- [x] Update keccak class with new design
- [x] Update keccak test
- [x] Check keccak correctness
- [x] Update tree builder rust wrappers
- [x] Leave doc comments

Future work:  
- [ ] Add keccak merkle tree builder externs
- [ ] Add keccak rust tree builder wrappers
- [ ] Write docs
- [ ] Add example
- [ ] Fix device output for tree builder

---------

Co-authored-by: Jeremy Felder <jeremy.felder1@gmail.com>
Co-authored-by: nonam3e <71525212+nonam3e@users.noreply.github.com>

wrappers/golang/core/poseidon.go

@@ -1,94 +0,0 @@
-package core
-
-import (
-	"fmt"
-	"unsafe"
-
-	cr "github.com/ingonyama-zk/icicle/v2/wrappers/golang/cuda_runtime"
-)
-
-type PoseidonConfig struct {
-	/// Details related to the device such as its id and stream id. See [DeviceContext](@ref device_context::DeviceContext).
-	Ctx                cr.DeviceContext
-	areInputsOnDevice  bool
-	areOutputsOnDevice bool
-	///If true, input is considered to be a states vector, holding the preimages in aligned or not aligned format.
-	///Memory under the input pointer will be used for states. If false, fresh states memory will be allocated and input will be copied into it */
-	InputIsAState bool
-	/// If true - input should be already aligned for poseidon permutation.
-	///* Aligned format: [0, A, B, 0, C, D, ...] (as you might get by using loop_state)
-	///* not aligned format: [A, B, 0, C, D, 0, ...] (as you might get from cudaMemcpy2D) */
-	Aligned bool
-	///If true, hash results will also be copied in the input pointer in aligned format
-	LoopState bool
-	///Whether to run the Poseidon asynchronously. If set to `true`, the poseidon_hash function will be
-	///non-blocking and you'd need to synchronize it explicitly by running `cudaStreamSynchronize` or `cudaDeviceSynchronize`.
-	///If set to false, the poseidon_hash function will block the current CPU thread. */
-	IsAsync bool
-}
-
-type PoseidonConstants[T any] struct {
-	Arity           int32
-	PartialRounds   int32
-	FullRoundsHalf  int32
-	RoundConstants  unsafe.Pointer
-	MdsMatrix       unsafe.Pointer
-	NonSparseMatrix unsafe.Pointer
-	SparseMatrices  unsafe.Pointer
-	DomainTag       T
-}
-
-func GetDefaultPoseidonConfig() PoseidonConfig {
-	ctx, _ := cr.GetDefaultDeviceContext()
-	return PoseidonConfig{
-		ctx,   // Ctx
-		false, // areInputsOnDevice
-		false, // areOutputsOnDevice
-		false, // inputIsAState
-		false, // aligned
-		false, // loopState
-		false, // IsAsync
-	}
-}
-
-func PoseidonCheck[T any](input, output HostOrDeviceSlice, cfg *PoseidonConfig, constants *PoseidonConstants[T], numberOfStates int) (unsafe.Pointer, unsafe.Pointer, unsafe.Pointer) {
-	inputLen, outputLen := input.Len(), output.Len()
-	arity := int(constants.Arity)
-	expectedInputLen := arity * numberOfStates
-	if cfg.InputIsAState {
-		expectedInputLen += numberOfStates
-	}
-
-	if inputLen != expectedInputLen {
-		errorString := fmt.Sprintf(
-			"input is not the right length for the given parameters: %d, should be: %d",
-			inputLen,
-			arity*numberOfStates,
-		)
-		panic(errorString)
-	}
-
-	if outputLen != numberOfStates {
-		errorString := fmt.Sprintf(
-			"output is not the right length for the given parameters: %d, should be: %d",
-			outputLen,
-			numberOfStates,
-		)
-		panic(errorString)
-	}
-	cfg.areInputsOnDevice = input.IsOnDevice()
-	cfg.areOutputsOnDevice = output.IsOnDevice()
-
-	if input.IsOnDevice() {
-		input.(DeviceSlice).CheckDevice()
-
-	}
-
-	if output.IsOnDevice() {
-		output.(DeviceSlice).CheckDevice()
-	}
-
-	cfgPointer := unsafe.Pointer(cfg)
-
-	return input.AsUnsafePointer(), output.AsUnsafePointer(), cfgPointer
-}

wrappers/golang/core/sponge.go

@@ -0,0 +1,105 @@
+package core
+
+import (
+	"fmt"
+
+	cr "github.com/ingonyama-zk/icicle/v2/wrappers/golang/cuda_runtime"
+)
+
+type SpongeConfig struct {
+	/// Details related to the device such as its id and stream.
+	Ctx cr.DeviceContext
+
+	areInputsOnDevice  bool
+	areResultsOnDevice bool
+
+	InputRate  uint32
+	OutputRate uint32
+	Offset     uint32
+
+	/// If true - input should be already aligned for poseidon permutation.
+	/// Aligned format: [0, A, B, 0, C, D, ...] (as you might get by using loop_state)
+	/// not aligned format: [A, B, 0, C, D, 0, ...] (as you might get from cudaMemcpy2D)
+	RecursiveSqueeze bool
+
+	/// If true, hash results will also be copied in the input pointer in aligned format
+	Aligned bool
+
+	/// Whether to run the SpongeHash asynchronously. If set to `true`, the SpongeHash function will be non-blocking
+	/// and you'd need to synchronize it explicitly by running `cudaStreamSynchronize` or `cudaDeviceSynchronize`.
+	/// If set to `false`, the SpongeHash function will block the current CPU thread.
+	IsAsync bool
+}
+
+func GetDefaultSpongeConfig() SpongeConfig {
+	ctx, _ := cr.GetDefaultDeviceContext()
+	return SpongeConfig{
+		ctx,
+		false,
+		false,
+		0,
+		0,
+		0,
+		false,
+		false,
+		false,
+	}
+}
+
+func SpongeInputCheck(inputs HostOrDeviceSlice, numberOfStates, inputBlockLength, inputRate uint32, ctx *cr.DeviceContext) {
+	if inputBlockLength > inputRate {
+		errorString := fmt.Sprintf(
+			"Input block (%d) can't be greater than input rate (%d)",
+			inputBlockLength,
+			inputRate,
+		)
+		panic(errorString)
+	}
+	inputsSizeExpected := inputBlockLength * numberOfStates
+	if inputs.Len() < int(inputsSizeExpected) {
+		errorString := fmt.Sprintf(
+			"inputs len is %d; but needs to be at least %d",
+			inputs.Len(),
+			inputsSizeExpected,
+		)
+		panic(errorString)
+	}
+	if inputs.IsOnDevice() {
+		inputs.(DeviceSlice).CheckDevice()
+	}
+}
+
+func SpongeStatesCheck(states DeviceSlice, numberOfStates, width uint32, ctx *cr.DeviceContext) {
+
+	statesSizeExpected := width * numberOfStates
+	if states.Len() < int(statesSizeExpected) {
+		errorString := fmt.Sprintf(
+			"inputs len is %d; but needs to be at least %d",
+			states.Len(),
+			statesSizeExpected,
+		)
+		panic(errorString)
+	}
+	states.CheckDevice()
+}
+
+func SpongeOutputsCheck(outputs HostOrDeviceSlice, numberOfStates, outputLen, width uint32, recursive bool, ctx *cr.DeviceContext) {
+	var outputsSizeExpected uint32
+	if recursive {
+		outputsSizeExpected = width * numberOfStates
+	} else {
+		outputsSizeExpected = outputLen * numberOfStates
+	}
+
+	if outputs.Len() < int(outputsSizeExpected) {
+		errorString := fmt.Sprintf(
+			"outputs len is %d; but needs to be at least %d",
+			outputs.Len(),
+			outputsSizeExpected,
+		)
+		panic(errorString)
+	}
+	if outputs.IsOnDevice() {
+		outputs.(DeviceSlice).CheckDevice()
+	}
+}

wrappers/golang/curves/bls12377/poseidon/include/poseidon.h

@@ -9,14 +9,40 @@ extern "C" {
 #endif
 
 typedef struct scalar_t scalar_t;
-typedef struct PoseidonConfig PoseidonConfig;
 typedef struct DeviceContext DeviceContext;
-typedef struct PoseidonConstants PoseidonConstants;
+typedef struct TreeBuilderConfig TreeBuilderConfig;
+typedef struct PoseidonInst PoseidonInst;
+typedef struct SpongeConfig SpongeConfig;
 
 
-cudaError_t bls12_377_poseidon_hash_cuda(const scalar_t* input, scalar_t* output, int number_of_states, int arity, PoseidonConstants* constants, PoseidonConfig* config);
-cudaError_t bls12_377_create_optimized_poseidon_constants_cuda(int arity, int full_rounds_halfs, int partial_rounds, const scalar_t* constants, DeviceContext* ctx, PoseidonConstants* poseidon_constants);
-cudaError_t bls12_377_init_optimized_poseidon_constants_cuda(int arity, DeviceContext* ctx, PoseidonConstants* constants);
+cudaError_t bls12_377_poseidon_create_cuda(
+  PoseidonInst** poseidon,
+  unsigned int arity,
+  unsigned int alpha,
+  unsigned int partial_rounds,
+  unsigned int full_rounds_half,
+  const scalar_t* round_constants,
+  const scalar_t* mds_matrix,
+  const scalar_t* non_sparse_matrix,
+  const scalar_t* sparse_matrices,
+  const scalar_t* domain_tag,
+  DeviceContext* ctx);
+
+cudaError_t bls12_377_poseidon_load_cuda(
+  PoseidonInst** poseidon,
+  unsigned int arity,
+  DeviceContext* ctx);
+
+cudaError_t bls12_377_poseidon_hash_many_cuda(
+  const PoseidonInst* poseidon,
+  const scalar_t* inputs,
+  scalar_t* output,
+  unsigned int number_of_states,
+  unsigned int input_block_len,
+  unsigned int output_len,
+  SpongeConfig* cfg);
+
+cudaError_t bls12_377_poseidon_delete_cuda(PoseidonInst* poseidon);
 
 #ifdef __cplusplus
 }

wrappers/golang/curves/bls12377/poseidon/poseidon.go

@@ -3,55 +3,85 @@ package poseidon
 // #cgo CFLAGS: -I./include/
 // #include "poseidon.h"
 import "C"
-
 import (
+	"runtime"
 	"unsafe"
 
 	"github.com/ingonyama-zk/icicle/v2/wrappers/golang/core"
 	cr "github.com/ingonyama-zk/icicle/v2/wrappers/golang/cuda_runtime"
+	bls12_377 "github.com/ingonyama-zk/icicle/v2/wrappers/golang/curves/bls12377"
 )
 
-func GetDefaultPoseidonConfig() core.PoseidonConfig {
-	return core.GetDefaultPoseidonConfig()
+type PoseidonHandler = C.struct_PoseidonInst
+type Poseidon struct {
+	width  uint32
+	handle *PoseidonHandler
 }
 
-func PoseidonHash[T any](scalars, results core.HostOrDeviceSlice, numberOfStates int, cfg *core.PoseidonConfig, constants *core.PoseidonConstants[T]) core.IcicleError {
-	scalarsPointer, resultsPointer, cfgPointer := core.PoseidonCheck(scalars, results, cfg, constants, numberOfStates)
+func Create(arity uint32, alpha uint32, fullRoundsHalf uint32, partialRounds uint32, scalars core.HostOrDeviceSlice, mdsMatrix core.HostOrDeviceSlice, nonSparseMatrix core.HostOrDeviceSlice, sparseMatrices core.HostOrDeviceSlice, domainTag bls12_377.ScalarField, ctx *cr.DeviceContext) (*Poseidon, core.IcicleError) {
+	var poseidon *PoseidonHandler
+	cArity := (C.uint)(arity)
+	cAlpha := (C.uint)(alpha)
+	cFullRoundsHalf := (C.uint)(fullRoundsHalf)
+	cPartialRounds := (C.uint)(partialRounds)
+	cScalars := (*C.scalar_t)(scalars.AsUnsafePointer())
+	cMdsMatrix := (*C.scalar_t)(mdsMatrix.AsUnsafePointer())
+	cNonSparseMatrix := (*C.scalar_t)(nonSparseMatrix.AsUnsafePointer())
+	cSparseMatrices := (*C.scalar_t)(sparseMatrices.AsUnsafePointer())
+	cDomainTag := (*C.scalar_t)(unsafe.Pointer(&domainTag))
+	cCtx := (*C.DeviceContext)(unsafe.Pointer(ctx))
+	__ret := C.bls12_377_poseidon_create_cuda(&poseidon, cArity, cAlpha, cFullRoundsHalf, cPartialRounds, cScalars, cMdsMatrix, cNonSparseMatrix, cSparseMatrices, cDomainTag, cCtx)
+	err := core.FromCudaError((cr.CudaError)(__ret))
+	if err.IcicleErrorCode != core.IcicleSuccess {
+		return nil, err
+	}
+	p := Poseidon{handle: poseidon, width: arity + 1}
+	runtime.SetFinalizer(&p, func(p *Poseidon) {
+		p.Delete()
+	})
+	return &p, err
+}
 
-	cScalars := (*C.scalar_t)(scalarsPointer)
-	cResults := (*C.scalar_t)(resultsPointer)
-	cNumberOfStates := (C.int)(numberOfStates)
-	cArity := (C.int)(constants.Arity)
-	cConstants := (*C.PoseidonConstants)(unsafe.Pointer(constants))
-	cCfg := (*C.PoseidonConfig)(cfgPointer)
+func Load(arity uint32, ctx *cr.DeviceContext) (*Poseidon, core.IcicleError) {
+	var poseidon *PoseidonHandler
+	cArity := (C.uint)(arity)
+	cCtx := (*C.DeviceContext)(unsafe.Pointer(ctx))
+	__ret := C.bls12_377_poseidon_load_cuda(&poseidon, cArity, cCtx)
+	err := core.FromCudaError((cr.CudaError)(__ret))
+	if err.IcicleErrorCode != core.IcicleSuccess {
+		return nil, err
+	}
+	p := Poseidon{handle: poseidon, width: arity + 1}
+	runtime.SetFinalizer(&p, func(p *Poseidon) {
+		p.Delete()
+	})
+	return &p, err
+}
 
-	__ret := C.bls12_377_poseidon_hash_cuda(cScalars, cResults, cNumberOfStates, cArity, cConstants, cCfg)
+func (poseidon *Poseidon) HashMany(inputs core.HostOrDeviceSlice, output core.HostOrDeviceSlice, numberOfStates uint32, inputBlockLen uint32, outputLen uint32, cfg *core.SpongeConfig) core.IcicleError {
+	core.SpongeInputCheck(inputs, numberOfStates, inputBlockLen, cfg.InputRate, &cfg.Ctx)
+	core.SpongeOutputsCheck(output, numberOfStates, outputLen, poseidon.width, false, &cfg.Ctx)
 
+	cInputs := (*C.scalar_t)(inputs.AsUnsafePointer())
+	cOutput := (*C.scalar_t)(output.AsUnsafePointer())
+	cNumberOfStates := (C.uint)(numberOfStates)
+	cInputBlockLen := (C.uint)(inputBlockLen)
+	cOutputLen := (C.uint)(outputLen)
+	cCfg := (*C.SpongeConfig)(unsafe.Pointer(cfg))
+	__ret := C.bls12_377_poseidon_hash_many_cuda(poseidon.handle, cInputs, cOutput, cNumberOfStates, cInputBlockLen, cOutputLen, cCfg)
 	err := (cr.CudaError)(__ret)
 	return core.FromCudaError(err)
 }
 
-func CreateOptimizedPoseidonConstants[T any](arity, fullRoundsHalfs, partialRounds int, constants core.HostOrDeviceSlice, ctx cr.DeviceContext, poseidonConstants *core.PoseidonConstants[T]) core.IcicleError {
-
-	cArity := (C.int)(arity)
-	cFullRoundsHalfs := (C.int)(fullRoundsHalfs)
-	cPartialRounds := (C.int)(partialRounds)
-	cConstants := (*C.scalar_t)(constants.AsUnsafePointer())
-	cCtx := (*C.DeviceContext)(unsafe.Pointer(&ctx))
-	cPoseidonConstants := (*C.PoseidonConstants)(unsafe.Pointer(poseidonConstants))
-
-	__ret := C.bls12_377_create_optimized_poseidon_constants_cuda(cArity, cFullRoundsHalfs, cPartialRounds, cConstants, cCtx, cPoseidonConstants)
+func (poseidon *Poseidon) Delete() core.IcicleError {
+	__ret := C.bls12_377_poseidon_delete_cuda(poseidon.handle)
 	err := (cr.CudaError)(__ret)
 	return core.FromCudaError(err)
 }
 
-func InitOptimizedPoseidonConstantsCuda[T any](arity int, ctx cr.DeviceContext, constants *core.PoseidonConstants[T]) core.IcicleError {
-
-	cArity := (C.int)(arity)
-	cCtx := (*C.DeviceContext)(unsafe.Pointer(&ctx))
-	cConstants := (*C.PoseidonConstants)(unsafe.Pointer(constants))
-
-	__ret := C.bls12_377_init_optimized_poseidon_constants_cuda(cArity, cCtx, cConstants)
-	err := (cr.CudaError)(__ret)
-	return core.FromCudaError(err)
+func (poseidon *Poseidon) GetDefaultSpongeConfig() core.SpongeConfig {
+	cfg := core.GetDefaultSpongeConfig()
+	cfg.InputRate = poseidon.width - 1
+	cfg.OutputRate = poseidon.width
+	return cfg
 }

wrappers/golang/curves/bls12377/tests/poseidon_test.go

@@ -7,6 +7,7 @@ import (
 	cr "github.com/ingonyama-zk/icicle/v2/wrappers/golang/cuda_runtime"
 	bls12_377 "github.com/ingonyama-zk/icicle/v2/wrappers/golang/curves/bls12377"
 	poseidon "github.com/ingonyama-zk/icicle/v2/wrappers/golang/curves/bls12377/poseidon"
+	"github.com/stretchr/testify/assert"
 )
 
 func TestPoseidon(t *testing.T) {
@@ -14,14 +15,11 @@ func TestPoseidon(t *testing.T) {
 	arity := 2
 	numberOfStates := 1
 
-	cfg := poseidon.GetDefaultPoseidonConfig()
-	cfg.IsAsync = true
-	stream, _ := cr.CreateStream()
-	cfg.Ctx.Stream = &stream
+	ctx, _ := cr.GetDefaultDeviceContext()
+	p, err := poseidon.Load(uint32(arity), &ctx)
+	assert.Equal(t, core.IcicleSuccess, err.IcicleErrorCode)
 
-	var constants core.PoseidonConstants[bls12_377.ScalarField]
-
-	poseidon.InitOptimizedPoseidonConstantsCuda(arity, cfg.Ctx, &constants) //generate constants
+	cfg := p.GetDefaultSpongeConfig()
 
 	scalars := bls12_377.GenerateScalars(numberOfStates * arity)
 	scalars[0] = scalars[0].Zero()
@@ -30,13 +28,13 @@ func TestPoseidon(t *testing.T) {
 	scalarsCopy := core.HostSliceFromElements(scalars[:numberOfStates*arity])
 
 	var deviceInput core.DeviceSlice
-	scalarsCopy.CopyToDeviceAsync(&deviceInput, stream, true)
+	scalarsCopy.CopyToDevice(&deviceInput, true)
 	var deviceOutput core.DeviceSlice
-	deviceOutput.MallocAsync(numberOfStates*scalarsCopy.SizeOfElement(), scalarsCopy.SizeOfElement(), stream)
+	deviceOutput.Malloc(numberOfStates*scalarsCopy.SizeOfElement(), scalarsCopy.SizeOfElement())
 
-	poseidon.PoseidonHash(deviceInput, deviceOutput, numberOfStates, &cfg, &constants) //run Hash function
+	err = p.HashMany(deviceInput, deviceOutput, uint32(numberOfStates), 1, 1, &cfg) //run Hash function
+	assert.Equal(t, core.IcicleSuccess, err.IcicleErrorCode)
 
 	output := make(core.HostSlice[bls12_377.ScalarField], numberOfStates)
-	output.CopyFromDeviceAsync(&deviceOutput, stream)
-
+	output.CopyFromDevice(&deviceOutput)
 }

wrappers/golang/curves/bls12381/poseidon/include/poseidon.h

@@ -9,14 +9,40 @@ extern "C" {
 #endif
 
 typedef struct scalar_t scalar_t;
-typedef struct PoseidonConfig PoseidonConfig;
 typedef struct DeviceContext DeviceContext;
-typedef struct PoseidonConstants PoseidonConstants;
+typedef struct TreeBuilderConfig TreeBuilderConfig;
+typedef struct PoseidonInst PoseidonInst;
+typedef struct SpongeConfig SpongeConfig;
 
 
-cudaError_t bls12_381_poseidon_hash_cuda(const scalar_t* input, scalar_t* output, int number_of_states, int arity, PoseidonConstants* constants, PoseidonConfig* config);
-cudaError_t bls12_381_create_optimized_poseidon_constants_cuda(int arity, int full_rounds_halfs, int partial_rounds, const scalar_t* constants, DeviceContext* ctx, PoseidonConstants* poseidon_constants);
-cudaError_t bls12_381_init_optimized_poseidon_constants_cuda(int arity, DeviceContext* ctx, PoseidonConstants* constants);
+cudaError_t bls12_381_poseidon_create_cuda(
+  PoseidonInst** poseidon,
+  unsigned int arity,
+  unsigned int alpha,
+  unsigned int partial_rounds,
+  unsigned int full_rounds_half,
+  const scalar_t* round_constants,
+  const scalar_t* mds_matrix,
+  const scalar_t* non_sparse_matrix,
+  const scalar_t* sparse_matrices,
+  const scalar_t* domain_tag,
+  DeviceContext* ctx);
+
+cudaError_t bls12_381_poseidon_load_cuda(
+  PoseidonInst** poseidon,
+  unsigned int arity,
+  DeviceContext* ctx);
+
+cudaError_t bls12_381_poseidon_hash_many_cuda(
+  const PoseidonInst* poseidon,
+  const scalar_t* inputs,
+  scalar_t* output,
+  unsigned int number_of_states,
+  unsigned int input_block_len,
+  unsigned int output_len,
+  SpongeConfig* cfg);
+
+cudaError_t bls12_381_poseidon_delete_cuda(PoseidonInst* poseidon);
 
 #ifdef __cplusplus
 }

wrappers/golang/curves/bls12381/poseidon/poseidon.go

@@ -3,55 +3,85 @@ package poseidon
 // #cgo CFLAGS: -I./include/
 // #include "poseidon.h"
 import "C"
-
 import (
+	"runtime"
 	"unsafe"
 
 	"github.com/ingonyama-zk/icicle/v2/wrappers/golang/core"
 	cr "github.com/ingonyama-zk/icicle/v2/wrappers/golang/cuda_runtime"
+	bls12_381 "github.com/ingonyama-zk/icicle/v2/wrappers/golang/curves/bls12381"
 )
 
-func GetDefaultPoseidonConfig() core.PoseidonConfig {
-	return core.GetDefaultPoseidonConfig()
+type PoseidonHandler = C.struct_PoseidonInst
+type Poseidon struct {
+	width  uint32
+	handle *PoseidonHandler
 }
 
-func PoseidonHash[T any](scalars, results core.HostOrDeviceSlice, numberOfStates int, cfg *core.PoseidonConfig, constants *core.PoseidonConstants[T]) core.IcicleError {
-	scalarsPointer, resultsPointer, cfgPointer := core.PoseidonCheck(scalars, results, cfg, constants, numberOfStates)
+func Create(arity uint32, alpha uint32, fullRoundsHalf uint32, partialRounds uint32, scalars core.HostOrDeviceSlice, mdsMatrix core.HostOrDeviceSlice, nonSparseMatrix core.HostOrDeviceSlice, sparseMatrices core.HostOrDeviceSlice, domainTag bls12_381.ScalarField, ctx *cr.DeviceContext) (*Poseidon, core.IcicleError) {
+	var poseidon *PoseidonHandler
+	cArity := (C.uint)(arity)
+	cAlpha := (C.uint)(alpha)
+	cFullRoundsHalf := (C.uint)(fullRoundsHalf)
+	cPartialRounds := (C.uint)(partialRounds)
+	cScalars := (*C.scalar_t)(scalars.AsUnsafePointer())
+	cMdsMatrix := (*C.scalar_t)(mdsMatrix.AsUnsafePointer())
+	cNonSparseMatrix := (*C.scalar_t)(nonSparseMatrix.AsUnsafePointer())
+	cSparseMatrices := (*C.scalar_t)(sparseMatrices.AsUnsafePointer())
+	cDomainTag := (*C.scalar_t)(unsafe.Pointer(&domainTag))
+	cCtx := (*C.DeviceContext)(unsafe.Pointer(ctx))
+	__ret := C.bls12_381_poseidon_create_cuda(&poseidon, cArity, cAlpha, cFullRoundsHalf, cPartialRounds, cScalars, cMdsMatrix, cNonSparseMatrix, cSparseMatrices, cDomainTag, cCtx)
+	err := core.FromCudaError((cr.CudaError)(__ret))
+	if err.IcicleErrorCode != core.IcicleSuccess {
+		return nil, err
+	}
+	p := Poseidon{handle: poseidon, width: arity + 1}
+	runtime.SetFinalizer(&p, func(p *Poseidon) {
+		p.Delete()
+	})
+	return &p, err
+}
 
-	cScalars := (*C.scalar_t)(scalarsPointer)
-	cResults := (*C.scalar_t)(resultsPointer)
-	cNumberOfStates := (C.int)(numberOfStates)
-	cArity := (C.int)(constants.Arity)
-	cConstants := (*C.PoseidonConstants)(unsafe.Pointer(constants))
-	cCfg := (*C.PoseidonConfig)(cfgPointer)
+func Load(arity uint32, ctx *cr.DeviceContext) (*Poseidon, core.IcicleError) {
+	var poseidon *PoseidonHandler
+	cArity := (C.uint)(arity)
+	cCtx := (*C.DeviceContext)(unsafe.Pointer(ctx))
+	__ret := C.bls12_381_poseidon_load_cuda(&poseidon, cArity, cCtx)
+	err := core.FromCudaError((cr.CudaError)(__ret))
+	if err.IcicleErrorCode != core.IcicleSuccess {
+		return nil, err
+	}
+	p := Poseidon{handle: poseidon, width: arity + 1}
+	runtime.SetFinalizer(&p, func(p *Poseidon) {
+		p.Delete()
+	})
+	return &p, err
+}
 
-	__ret := C.bls12_381_poseidon_hash_cuda(cScalars, cResults, cNumberOfStates, cArity, cConstants, cCfg)
+func (poseidon *Poseidon) HashMany(inputs core.HostOrDeviceSlice, output core.HostOrDeviceSlice, numberOfStates uint32, inputBlockLen uint32, outputLen uint32, cfg *core.SpongeConfig) core.IcicleError {
+	core.SpongeInputCheck(inputs, numberOfStates, inputBlockLen, cfg.InputRate, &cfg.Ctx)
+	core.SpongeOutputsCheck(output, numberOfStates, outputLen, poseidon.width, false, &cfg.Ctx)
 
+	cInputs := (*C.scalar_t)(inputs.AsUnsafePointer())
+	cOutput := (*C.scalar_t)(output.AsUnsafePointer())
+	cNumberOfStates := (C.uint)(numberOfStates)
+	cInputBlockLen := (C.uint)(inputBlockLen)
+	cOutputLen := (C.uint)(outputLen)
+	cCfg := (*C.SpongeConfig)(unsafe.Pointer(cfg))
+	__ret := C.bls12_381_poseidon_hash_many_cuda(poseidon.handle, cInputs, cOutput, cNumberOfStates, cInputBlockLen, cOutputLen, cCfg)
 	err := (cr.CudaError)(__ret)
 	return core.FromCudaError(err)
 }
 
-func CreateOptimizedPoseidonConstants[T any](arity, fullRoundsHalfs, partialRounds int, constants core.HostOrDeviceSlice, ctx cr.DeviceContext, poseidonConstants *core.PoseidonConstants[T]) core.IcicleError {
-
-	cArity := (C.int)(arity)
-	cFullRoundsHalfs := (C.int)(fullRoundsHalfs)
-	cPartialRounds := (C.int)(partialRounds)
-	cConstants := (*C.scalar_t)(constants.AsUnsafePointer())
-	cCtx := (*C.DeviceContext)(unsafe.Pointer(&ctx))
-	cPoseidonConstants := (*C.PoseidonConstants)(unsafe.Pointer(poseidonConstants))
-
-	__ret := C.bls12_381_create_optimized_poseidon_constants_cuda(cArity, cFullRoundsHalfs, cPartialRounds, cConstants, cCtx, cPoseidonConstants)
+func (poseidon *Poseidon) Delete() core.IcicleError {
+	__ret := C.bls12_381_poseidon_delete_cuda(poseidon.handle)
 	err := (cr.CudaError)(__ret)
 	return core.FromCudaError(err)
 }
 
-func InitOptimizedPoseidonConstantsCuda[T any](arity int, ctx cr.DeviceContext, constants *core.PoseidonConstants[T]) core.IcicleError {
-
-	cArity := (C.int)(arity)
-	cCtx := (*C.DeviceContext)(unsafe.Pointer(&ctx))
-	cConstants := (*C.PoseidonConstants)(unsafe.Pointer(constants))
-
-	__ret := C.bls12_381_init_optimized_poseidon_constants_cuda(cArity, cCtx, cConstants)
-	err := (cr.CudaError)(__ret)
-	return core.FromCudaError(err)
+func (poseidon *Poseidon) GetDefaultSpongeConfig() core.SpongeConfig {
+	cfg := core.GetDefaultSpongeConfig()
+	cfg.InputRate = poseidon.width - 1
+	cfg.OutputRate = poseidon.width
+	return cfg
 }

wrappers/golang/curves/bls12381/tests/poseidon_test.go

@@ -7,29 +7,19 @@ import (
 	cr "github.com/ingonyama-zk/icicle/v2/wrappers/golang/cuda_runtime"
 	bls12_381 "github.com/ingonyama-zk/icicle/v2/wrappers/golang/curves/bls12381"
 	poseidon "github.com/ingonyama-zk/icicle/v2/wrappers/golang/curves/bls12381/poseidon"
-
-	"fmt"
 	"github.com/stretchr/testify/assert"
 )
 
-func formatOutput(x bls12_381.ScalarField) string {
-	r := x.GetLimbs()
-	return fmt.Sprintf("%08x%08x%08x%08x%08x%08x%08x%08x", r[7], r[6], r[5], r[4], r[3], r[2], r[1], r[0])
-}
-
 func TestPoseidon(t *testing.T) {
 
 	arity := 2
 	numberOfStates := 1
 
-	cfg := poseidon.GetDefaultPoseidonConfig()
-	cfg.IsAsync = true
-	stream, _ := cr.CreateStream()
-	cfg.Ctx.Stream = &stream
+	ctx, _ := cr.GetDefaultDeviceContext()
+	p, err := poseidon.Load(uint32(arity), &ctx)
+	assert.Equal(t, core.IcicleSuccess, err.IcicleErrorCode)
 
-	var constants core.PoseidonConstants[bls12_381.ScalarField]
-
-	poseidon.InitOptimizedPoseidonConstantsCuda(arity, cfg.Ctx, &constants) //generate constants
+	cfg := p.GetDefaultSpongeConfig()
 
 	scalars := bls12_381.GenerateScalars(numberOfStates * arity)
 	scalars[0] = scalars[0].Zero()
@@ -38,18 +28,13 @@ func TestPoseidon(t *testing.T) {
 	scalarsCopy := core.HostSliceFromElements(scalars[:numberOfStates*arity])
 
 	var deviceInput core.DeviceSlice
-	scalarsCopy.CopyToDeviceAsync(&deviceInput, stream, true)
+	scalarsCopy.CopyToDevice(&deviceInput, true)
 	var deviceOutput core.DeviceSlice
-	deviceOutput.MallocAsync(numberOfStates*scalarsCopy.SizeOfElement(), scalarsCopy.SizeOfElement(), stream)
+	deviceOutput.Malloc(numberOfStates*scalarsCopy.SizeOfElement(), scalarsCopy.SizeOfElement())
 
-	poseidon.PoseidonHash(deviceInput, deviceOutput, numberOfStates, &cfg, &constants) //run Hash function
+	err = p.HashMany(deviceInput, deviceOutput, uint32(numberOfStates), 1, 1, &cfg) //run Hash function
+	assert.Equal(t, core.IcicleSuccess, err.IcicleErrorCode)
 
 	output := make(core.HostSlice[bls12_381.ScalarField], numberOfStates)
-	output.CopyFromDeviceAsync(&deviceOutput, stream)
-
-	expectedString := "48fe0b1331196f6cdb33a7c6e5af61b76fd388e1ef1d3d418be5147f0e4613d4" //This result is from https://github.com/triplewz/poseidon
-	outputString := formatOutput(output[0])
-
-	assert.Equal(t, outputString, expectedString, "Poseidon hash does not match expected result")
-
+	output.CopyFromDevice(&deviceOutput)
 }

wrappers/golang/curves/bn254/poseidon/include/poseidon.h

@@ -9,14 +9,40 @@ extern "C" {
 #endif
 
 typedef struct scalar_t scalar_t;
-typedef struct PoseidonConfig PoseidonConfig;
 typedef struct DeviceContext DeviceContext;
-typedef struct PoseidonConstants PoseidonConstants;
+typedef struct TreeBuilderConfig TreeBuilderConfig;
+typedef struct PoseidonInst PoseidonInst;
+typedef struct SpongeConfig SpongeConfig;
 
 
-cudaError_t bn254_poseidon_hash_cuda(const scalar_t* input, scalar_t* output, int number_of_states, int arity, PoseidonConstants* constants, PoseidonConfig* config);
-cudaError_t bn254_create_optimized_poseidon_constants_cuda(int arity, int full_rounds_halfs, int partial_rounds, const scalar_t* constants, DeviceContext* ctx, PoseidonConstants* poseidon_constants);
-cudaError_t bn254_init_optimized_poseidon_constants_cuda(int arity, DeviceContext* ctx, PoseidonConstants* constants);
+cudaError_t bn254_poseidon_create_cuda(
+  PoseidonInst** poseidon,
+  unsigned int arity,
+  unsigned int alpha,
+  unsigned int partial_rounds,
+  unsigned int full_rounds_half,
+  const scalar_t* round_constants,
+  const scalar_t* mds_matrix,
+  const scalar_t* non_sparse_matrix,
+  const scalar_t* sparse_matrices,
+  const scalar_t* domain_tag,
+  DeviceContext* ctx);
+
+cudaError_t bn254_poseidon_load_cuda(
+  PoseidonInst** poseidon,
+  unsigned int arity,
+  DeviceContext* ctx);
+
+cudaError_t bn254_poseidon_hash_many_cuda(
+  const PoseidonInst* poseidon,
+  const scalar_t* inputs,
+  scalar_t* output,
+  unsigned int number_of_states,
+  unsigned int input_block_len,
+  unsigned int output_len,
+  SpongeConfig* cfg);
+
+cudaError_t bn254_poseidon_delete_cuda(PoseidonInst* poseidon);
 
 #ifdef __cplusplus
 }

wrappers/golang/curves/bn254/poseidon/poseidon.go

@@ -3,55 +3,85 @@ package poseidon
 // #cgo CFLAGS: -I./include/
 // #include "poseidon.h"
 import "C"
-
 import (
+	"runtime"
 	"unsafe"
 
 	"github.com/ingonyama-zk/icicle/v2/wrappers/golang/core"
 	cr "github.com/ingonyama-zk/icicle/v2/wrappers/golang/cuda_runtime"
+	bn254 "github.com/ingonyama-zk/icicle/v2/wrappers/golang/curves/bn254"
 )
 
-func GetDefaultPoseidonConfig() core.PoseidonConfig {
-	return core.GetDefaultPoseidonConfig()
+type PoseidonHandler = C.struct_PoseidonInst
+type Poseidon struct {
+	width  uint32
+	handle *PoseidonHandler
 }
 
-func PoseidonHash[T any](scalars, results core.HostOrDeviceSlice, numberOfStates int, cfg *core.PoseidonConfig, constants *core.PoseidonConstants[T]) core.IcicleError {
-	scalarsPointer, resultsPointer, cfgPointer := core.PoseidonCheck(scalars, results, cfg, constants, numberOfStates)
+func Create(arity uint32, alpha uint32, fullRoundsHalf uint32, partialRounds uint32, scalars core.HostOrDeviceSlice, mdsMatrix core.HostOrDeviceSlice, nonSparseMatrix core.HostOrDeviceSlice, sparseMatrices core.HostOrDeviceSlice, domainTag bn254.ScalarField, ctx *cr.DeviceContext) (*Poseidon, core.IcicleError) {
+	var poseidon *PoseidonHandler
+	cArity := (C.uint)(arity)
+	cAlpha := (C.uint)(alpha)
+	cFullRoundsHalf := (C.uint)(fullRoundsHalf)
+	cPartialRounds := (C.uint)(partialRounds)
+	cScalars := (*C.scalar_t)(scalars.AsUnsafePointer())
+	cMdsMatrix := (*C.scalar_t)(mdsMatrix.AsUnsafePointer())
+	cNonSparseMatrix := (*C.scalar_t)(nonSparseMatrix.AsUnsafePointer())
+	cSparseMatrices := (*C.scalar_t)(sparseMatrices.AsUnsafePointer())
+	cDomainTag := (*C.scalar_t)(unsafe.Pointer(&domainTag))
+	cCtx := (*C.DeviceContext)(unsafe.Pointer(ctx))
+	__ret := C.bn254_poseidon_create_cuda(&poseidon, cArity, cAlpha, cFullRoundsHalf, cPartialRounds, cScalars, cMdsMatrix, cNonSparseMatrix, cSparseMatrices, cDomainTag, cCtx)
+	err := core.FromCudaError((cr.CudaError)(__ret))
+	if err.IcicleErrorCode != core.IcicleSuccess {
+		return nil, err
+	}
+	p := Poseidon{handle: poseidon, width: arity + 1}
+	runtime.SetFinalizer(&p, func(p *Poseidon) {
+		p.Delete()
+	})
+	return &p, err
+}
 
-	cScalars := (*C.scalar_t)(scalarsPointer)
-	cResults := (*C.scalar_t)(resultsPointer)
-	cNumberOfStates := (C.int)(numberOfStates)
-	cArity := (C.int)(constants.Arity)
-	cConstants := (*C.PoseidonConstants)(unsafe.Pointer(constants))
-	cCfg := (*C.PoseidonConfig)(cfgPointer)
+func Load(arity uint32, ctx *cr.DeviceContext) (*Poseidon, core.IcicleError) {
+	var poseidon *PoseidonHandler
+	cArity := (C.uint)(arity)
+	cCtx := (*C.DeviceContext)(unsafe.Pointer(ctx))
+	__ret := C.bn254_poseidon_load_cuda(&poseidon, cArity, cCtx)
+	err := core.FromCudaError((cr.CudaError)(__ret))
+	if err.IcicleErrorCode != core.IcicleSuccess {
+		return nil, err
+	}
+	p := Poseidon{handle: poseidon, width: arity + 1}
+	runtime.SetFinalizer(&p, func(p *Poseidon) {
+		p.Delete()
+	})
+	return &p, err
+}
 
-	__ret := C.bn254_poseidon_hash_cuda(cScalars, cResults, cNumberOfStates, cArity, cConstants, cCfg)
+func (poseidon *Poseidon) HashMany(inputs core.HostOrDeviceSlice, output core.HostOrDeviceSlice, numberOfStates uint32, inputBlockLen uint32, outputLen uint32, cfg *core.SpongeConfig) core.IcicleError {
+	core.SpongeInputCheck(inputs, numberOfStates, inputBlockLen, cfg.InputRate, &cfg.Ctx)
+	core.SpongeOutputsCheck(output, numberOfStates, outputLen, poseidon.width, false, &cfg.Ctx)
 
+	cInputs := (*C.scalar_t)(inputs.AsUnsafePointer())
+	cOutput := (*C.scalar_t)(output.AsUnsafePointer())
+	cNumberOfStates := (C.uint)(numberOfStates)
+	cInputBlockLen := (C.uint)(inputBlockLen)
+	cOutputLen := (C.uint)(outputLen)
+	cCfg := (*C.SpongeConfig)(unsafe.Pointer(cfg))
+	__ret := C.bn254_poseidon_hash_many_cuda(poseidon.handle, cInputs, cOutput, cNumberOfStates, cInputBlockLen, cOutputLen, cCfg)
 	err := (cr.CudaError)(__ret)
 	return core.FromCudaError(err)
 }
 
-func CreateOptimizedPoseidonConstants[T any](arity, fullRoundsHalfs, partialRounds int, constants core.HostOrDeviceSlice, ctx cr.DeviceContext, poseidonConstants *core.PoseidonConstants[T]) core.IcicleError {
-
-	cArity := (C.int)(arity)
-	cFullRoundsHalfs := (C.int)(fullRoundsHalfs)
-	cPartialRounds := (C.int)(partialRounds)
-	cConstants := (*C.scalar_t)(constants.AsUnsafePointer())
-	cCtx := (*C.DeviceContext)(unsafe.Pointer(&ctx))
-	cPoseidonConstants := (*C.PoseidonConstants)(unsafe.Pointer(poseidonConstants))
-
-	__ret := C.bn254_create_optimized_poseidon_constants_cuda(cArity, cFullRoundsHalfs, cPartialRounds, cConstants, cCtx, cPoseidonConstants)
+func (poseidon *Poseidon) Delete() core.IcicleError {
+	__ret := C.bn254_poseidon_delete_cuda(poseidon.handle)
 	err := (cr.CudaError)(__ret)
 	return core.FromCudaError(err)
 }
 
-func InitOptimizedPoseidonConstantsCuda[T any](arity int, ctx cr.DeviceContext, constants *core.PoseidonConstants[T]) core.IcicleError {
-
-	cArity := (C.int)(arity)
-	cCtx := (*C.DeviceContext)(unsafe.Pointer(&ctx))
-	cConstants := (*C.PoseidonConstants)(unsafe.Pointer(constants))
-
-	__ret := C.bn254_init_optimized_poseidon_constants_cuda(cArity, cCtx, cConstants)
-	err := (cr.CudaError)(__ret)
-	return core.FromCudaError(err)
+func (poseidon *Poseidon) GetDefaultSpongeConfig() core.SpongeConfig {
+	cfg := core.GetDefaultSpongeConfig()
+	cfg.InputRate = poseidon.width - 1
+	cfg.OutputRate = poseidon.width
+	return cfg
 }

wrappers/golang/curves/bn254/tests/poseidon_test.go

@@ -7,6 +7,7 @@ import (
 	cr "github.com/ingonyama-zk/icicle/v2/wrappers/golang/cuda_runtime"
 	bn254 "github.com/ingonyama-zk/icicle/v2/wrappers/golang/curves/bn254"
 	poseidon "github.com/ingonyama-zk/icicle/v2/wrappers/golang/curves/bn254/poseidon"
+	"github.com/stretchr/testify/assert"
 )
 
 func TestPoseidon(t *testing.T) {
@@ -14,14 +15,11 @@ func TestPoseidon(t *testing.T) {
 	arity := 2
 	numberOfStates := 1
 
-	cfg := poseidon.GetDefaultPoseidonConfig()
-	cfg.IsAsync = true
-	stream, _ := cr.CreateStream()
-	cfg.Ctx.Stream = &stream
+	ctx, _ := cr.GetDefaultDeviceContext()
+	p, err := poseidon.Load(uint32(arity), &ctx)
+	assert.Equal(t, core.IcicleSuccess, err.IcicleErrorCode)
 
-	var constants core.PoseidonConstants[bn254.ScalarField]
-
-	poseidon.InitOptimizedPoseidonConstantsCuda(arity, cfg.Ctx, &constants) //generate constants
+	cfg := p.GetDefaultSpongeConfig()
 
 	scalars := bn254.GenerateScalars(numberOfStates * arity)
 	scalars[0] = scalars[0].Zero()
@@ -30,13 +28,13 @@ func TestPoseidon(t *testing.T) {
 	scalarsCopy := core.HostSliceFromElements(scalars[:numberOfStates*arity])
 
 	var deviceInput core.DeviceSlice
-	scalarsCopy.CopyToDeviceAsync(&deviceInput, stream, true)
+	scalarsCopy.CopyToDevice(&deviceInput, true)
 	var deviceOutput core.DeviceSlice
-	deviceOutput.MallocAsync(numberOfStates*scalarsCopy.SizeOfElement(), scalarsCopy.SizeOfElement(), stream)
+	deviceOutput.Malloc(numberOfStates*scalarsCopy.SizeOfElement(), scalarsCopy.SizeOfElement())
 
-	poseidon.PoseidonHash(deviceInput, deviceOutput, numberOfStates, &cfg, &constants) //run Hash function
+	err = p.HashMany(deviceInput, deviceOutput, uint32(numberOfStates), 1, 1, &cfg) //run Hash function
+	assert.Equal(t, core.IcicleSuccess, err.IcicleErrorCode)
 
 	output := make(core.HostSlice[bn254.ScalarField], numberOfStates)
-	output.CopyFromDeviceAsync(&deviceOutput, stream)
-
+	output.CopyFromDevice(&deviceOutput)
 }

wrappers/golang/curves/bw6761/poseidon/include/poseidon.h

@@ -9,14 +9,40 @@ extern "C" {
 #endif
 
 typedef struct scalar_t scalar_t;
-typedef struct PoseidonConfig PoseidonConfig;
 typedef struct DeviceContext DeviceContext;
-typedef struct PoseidonConstants PoseidonConstants;
+typedef struct TreeBuilderConfig TreeBuilderConfig;
+typedef struct PoseidonInst PoseidonInst;
+typedef struct SpongeConfig SpongeConfig;
 
 
-cudaError_t bw6_761_poseidon_hash_cuda(const scalar_t* input, scalar_t* output, int number_of_states, int arity, PoseidonConstants* constants, PoseidonConfig* config);
-cudaError_t bw6_761_create_optimized_poseidon_constants_cuda(int arity, int full_rounds_halfs, int partial_rounds, const scalar_t* constants, DeviceContext* ctx, PoseidonConstants* poseidon_constants);
-cudaError_t bw6_761_init_optimized_poseidon_constants_cuda(int arity, DeviceContext* ctx, PoseidonConstants* constants);
+cudaError_t bw6_761_poseidon_create_cuda(
+  PoseidonInst** poseidon,
+  unsigned int arity,
+  unsigned int alpha,
+  unsigned int partial_rounds,
+  unsigned int full_rounds_half,
+  const scalar_t* round_constants,
+  const scalar_t* mds_matrix,
+  const scalar_t* non_sparse_matrix,
+  const scalar_t* sparse_matrices,
+  const scalar_t* domain_tag,
+  DeviceContext* ctx);
+
+cudaError_t bw6_761_poseidon_load_cuda(
+  PoseidonInst** poseidon,
+  unsigned int arity,
+  DeviceContext* ctx);
+
+cudaError_t bw6_761_poseidon_hash_many_cuda(
+  const PoseidonInst* poseidon,
+  const scalar_t* inputs,
+  scalar_t* output,
+  unsigned int number_of_states,
+  unsigned int input_block_len,
+  unsigned int output_len,
+  SpongeConfig* cfg);
+
+cudaError_t bw6_761_poseidon_delete_cuda(PoseidonInst* poseidon);
 
 #ifdef __cplusplus
 }

wrappers/golang/curves/bw6761/poseidon/poseidon.go

@@ -3,55 +3,85 @@ package poseidon
 // #cgo CFLAGS: -I./include/
 // #include "poseidon.h"
 import "C"
-
 import (
+	"runtime"
 	"unsafe"
 
 	"github.com/ingonyama-zk/icicle/v2/wrappers/golang/core"
 	cr "github.com/ingonyama-zk/icicle/v2/wrappers/golang/cuda_runtime"
+	bw6_761 "github.com/ingonyama-zk/icicle/v2/wrappers/golang/curves/bw6761"
 )
 
-func GetDefaultPoseidonConfig() core.PoseidonConfig {
-	return core.GetDefaultPoseidonConfig()
+type PoseidonHandler = C.struct_PoseidonInst
+type Poseidon struct {
+	width  uint32
+	handle *PoseidonHandler
 }
 
-func PoseidonHash[T any](scalars, results core.HostOrDeviceSlice, numberOfStates int, cfg *core.PoseidonConfig, constants *core.PoseidonConstants[T]) core.IcicleError {
-	scalarsPointer, resultsPointer, cfgPointer := core.PoseidonCheck(scalars, results, cfg, constants, numberOfStates)
+func Create(arity uint32, alpha uint32, fullRoundsHalf uint32, partialRounds uint32, scalars core.HostOrDeviceSlice, mdsMatrix core.HostOrDeviceSlice, nonSparseMatrix core.HostOrDeviceSlice, sparseMatrices core.HostOrDeviceSlice, domainTag bw6_761.ScalarField, ctx *cr.DeviceContext) (*Poseidon, core.IcicleError) {
+	var poseidon *PoseidonHandler
+	cArity := (C.uint)(arity)
+	cAlpha := (C.uint)(alpha)
+	cFullRoundsHalf := (C.uint)(fullRoundsHalf)
+	cPartialRounds := (C.uint)(partialRounds)
+	cScalars := (*C.scalar_t)(scalars.AsUnsafePointer())
+	cMdsMatrix := (*C.scalar_t)(mdsMatrix.AsUnsafePointer())
+	cNonSparseMatrix := (*C.scalar_t)(nonSparseMatrix.AsUnsafePointer())
+	cSparseMatrices := (*C.scalar_t)(sparseMatrices.AsUnsafePointer())
+	cDomainTag := (*C.scalar_t)(unsafe.Pointer(&domainTag))
+	cCtx := (*C.DeviceContext)(unsafe.Pointer(ctx))
+	__ret := C.bw6_761_poseidon_create_cuda(&poseidon, cArity, cAlpha, cFullRoundsHalf, cPartialRounds, cScalars, cMdsMatrix, cNonSparseMatrix, cSparseMatrices, cDomainTag, cCtx)
+	err := core.FromCudaError((cr.CudaError)(__ret))
+	if err.IcicleErrorCode != core.IcicleSuccess {
+		return nil, err
+	}
+	p := Poseidon{handle: poseidon, width: arity + 1}
+	runtime.SetFinalizer(&p, func(p *Poseidon) {
+		p.Delete()
+	})
+	return &p, err
+}
 
-	cScalars := (*C.scalar_t)(scalarsPointer)
-	cResults := (*C.scalar_t)(resultsPointer)
-	cNumberOfStates := (C.int)(numberOfStates)
-	cArity := (C.int)(constants.Arity)
-	cConstants := (*C.PoseidonConstants)(unsafe.Pointer(constants))
-	cCfg := (*C.PoseidonConfig)(cfgPointer)
+func Load(arity uint32, ctx *cr.DeviceContext) (*Poseidon, core.IcicleError) {
+	var poseidon *PoseidonHandler
+	cArity := (C.uint)(arity)
+	cCtx := (*C.DeviceContext)(unsafe.Pointer(ctx))
+	__ret := C.bw6_761_poseidon_load_cuda(&poseidon, cArity, cCtx)
+	err := core.FromCudaError((cr.CudaError)(__ret))
+	if err.IcicleErrorCode != core.IcicleSuccess {
+		return nil, err
+	}
+	p := Poseidon{handle: poseidon, width: arity + 1}
+	runtime.SetFinalizer(&p, func(p *Poseidon) {
+		p.Delete()
+	})
+	return &p, err
+}
 
-	__ret := C.bw6_761_poseidon_hash_cuda(cScalars, cResults, cNumberOfStates, cArity, cConstants, cCfg)
+func (poseidon *Poseidon) HashMany(inputs core.HostOrDeviceSlice, output core.HostOrDeviceSlice, numberOfStates uint32, inputBlockLen uint32, outputLen uint32, cfg *core.SpongeConfig) core.IcicleError {
+	core.SpongeInputCheck(inputs, numberOfStates, inputBlockLen, cfg.InputRate, &cfg.Ctx)
+	core.SpongeOutputsCheck(output, numberOfStates, outputLen, poseidon.width, false, &cfg.Ctx)
 
+	cInputs := (*C.scalar_t)(inputs.AsUnsafePointer())
+	cOutput := (*C.scalar_t)(output.AsUnsafePointer())
+	cNumberOfStates := (C.uint)(numberOfStates)
+	cInputBlockLen := (C.uint)(inputBlockLen)
+	cOutputLen := (C.uint)(outputLen)
+	cCfg := (*C.SpongeConfig)(unsafe.Pointer(cfg))
+	__ret := C.bw6_761_poseidon_hash_many_cuda(poseidon.handle, cInputs, cOutput, cNumberOfStates, cInputBlockLen, cOutputLen, cCfg)
 	err := (cr.CudaError)(__ret)
 	return core.FromCudaError(err)
 }
 
-func CreateOptimizedPoseidonConstants[T any](arity, fullRoundsHalfs, partialRounds int, constants core.HostOrDeviceSlice, ctx cr.DeviceContext, poseidonConstants *core.PoseidonConstants[T]) core.IcicleError {
-
-	cArity := (C.int)(arity)
-	cFullRoundsHalfs := (C.int)(fullRoundsHalfs)
-	cPartialRounds := (C.int)(partialRounds)
-	cConstants := (*C.scalar_t)(constants.AsUnsafePointer())
-	cCtx := (*C.DeviceContext)(unsafe.Pointer(&ctx))
-	cPoseidonConstants := (*C.PoseidonConstants)(unsafe.Pointer(poseidonConstants))
-
-	__ret := C.bw6_761_create_optimized_poseidon_constants_cuda(cArity, cFullRoundsHalfs, cPartialRounds, cConstants, cCtx, cPoseidonConstants)
+func (poseidon *Poseidon) Delete() core.IcicleError {
+	__ret := C.bw6_761_poseidon_delete_cuda(poseidon.handle)
 	err := (cr.CudaError)(__ret)
 	return core.FromCudaError(err)
 }
 
-func InitOptimizedPoseidonConstantsCuda[T any](arity int, ctx cr.DeviceContext, constants *core.PoseidonConstants[T]) core.IcicleError {
-
-	cArity := (C.int)(arity)
-	cCtx := (*C.DeviceContext)(unsafe.Pointer(&ctx))
-	cConstants := (*C.PoseidonConstants)(unsafe.Pointer(constants))
-
-	__ret := C.bw6_761_init_optimized_poseidon_constants_cuda(cArity, cCtx, cConstants)
-	err := (cr.CudaError)(__ret)
-	return core.FromCudaError(err)
+func (poseidon *Poseidon) GetDefaultSpongeConfig() core.SpongeConfig {
+	cfg := core.GetDefaultSpongeConfig()
+	cfg.InputRate = poseidon.width - 1
+	cfg.OutputRate = poseidon.width
+	return cfg
 }

wrappers/golang/curves/bw6761/tests/poseidon_test.go

@@ -7,6 +7,7 @@ import (
 	cr "github.com/ingonyama-zk/icicle/v2/wrappers/golang/cuda_runtime"
 	bw6_761 "github.com/ingonyama-zk/icicle/v2/wrappers/golang/curves/bw6761"
 	poseidon "github.com/ingonyama-zk/icicle/v2/wrappers/golang/curves/bw6761/poseidon"
+	"github.com/stretchr/testify/assert"
 )
 
 func TestPoseidon(t *testing.T) {
@@ -14,14 +15,11 @@ func TestPoseidon(t *testing.T) {
 	arity := 2
 	numberOfStates := 1
 
-	cfg := poseidon.GetDefaultPoseidonConfig()
-	cfg.IsAsync = true
-	stream, _ := cr.CreateStream()
-	cfg.Ctx.Stream = &stream
+	ctx, _ := cr.GetDefaultDeviceContext()
+	p, err := poseidon.Load(uint32(arity), &ctx)
+	assert.Equal(t, core.IcicleSuccess, err.IcicleErrorCode)
 
-	var constants core.PoseidonConstants[bw6_761.ScalarField]
-
-	poseidon.InitOptimizedPoseidonConstantsCuda(arity, cfg.Ctx, &constants) //generate constants
+	cfg := p.GetDefaultSpongeConfig()
 
 	scalars := bw6_761.GenerateScalars(numberOfStates * arity)
 	scalars[0] = scalars[0].Zero()
@@ -30,13 +28,13 @@ func TestPoseidon(t *testing.T) {
 	scalarsCopy := core.HostSliceFromElements(scalars[:numberOfStates*arity])
 
 	var deviceInput core.DeviceSlice
-	scalarsCopy.CopyToDeviceAsync(&deviceInput, stream, true)
+	scalarsCopy.CopyToDevice(&deviceInput, true)
 	var deviceOutput core.DeviceSlice
-	deviceOutput.MallocAsync(numberOfStates*scalarsCopy.SizeOfElement(), scalarsCopy.SizeOfElement(), stream)
+	deviceOutput.Malloc(numberOfStates*scalarsCopy.SizeOfElement(), scalarsCopy.SizeOfElement())
 
-	poseidon.PoseidonHash(deviceInput, deviceOutput, numberOfStates, &cfg, &constants) //run Hash function
+	err = p.HashMany(deviceInput, deviceOutput, uint32(numberOfStates), 1, 1, &cfg) //run Hash function
+	assert.Equal(t, core.IcicleSuccess, err.IcicleErrorCode)
 
 	output := make(core.HostSlice[bw6_761.ScalarField], numberOfStates)
-	output.CopyFromDeviceAsync(&deviceOutput, stream)
-
+	output.CopyFromDevice(&deviceOutput)
 }

wrappers/golang/curves/grumpkin/poseidon/include/poseidon.h

@@ -9,14 +9,40 @@ extern "C" {
 #endif
 
 typedef struct scalar_t scalar_t;
-typedef struct PoseidonConfig PoseidonConfig;
 typedef struct DeviceContext DeviceContext;
-typedef struct PoseidonConstants PoseidonConstants;
+typedef struct TreeBuilderConfig TreeBuilderConfig;
+typedef struct PoseidonInst PoseidonInst;
+typedef struct SpongeConfig SpongeConfig;
 
 
-cudaError_t grumpkin_poseidon_hash_cuda(const scalar_t* input, scalar_t* output, int number_of_states, int arity, PoseidonConstants* constants, PoseidonConfig* config);
-cudaError_t grumpkin_create_optimized_poseidon_constants_cuda(int arity, int full_rounds_halfs, int partial_rounds, const scalar_t* constants, DeviceContext* ctx, PoseidonConstants* poseidon_constants);
-cudaError_t grumpkin_init_optimized_poseidon_constants_cuda(int arity, DeviceContext* ctx, PoseidonConstants* constants);
+cudaError_t grumpkin_poseidon_create_cuda(
+  PoseidonInst** poseidon,
+  unsigned int arity,
+  unsigned int alpha,
+  unsigned int partial_rounds,
+  unsigned int full_rounds_half,
+  const scalar_t* round_constants,
+  const scalar_t* mds_matrix,
+  const scalar_t* non_sparse_matrix,
+  const scalar_t* sparse_matrices,
+  const scalar_t* domain_tag,
+  DeviceContext* ctx);
+
+cudaError_t grumpkin_poseidon_load_cuda(
+  PoseidonInst** poseidon,
+  unsigned int arity,
+  DeviceContext* ctx);
+
+cudaError_t grumpkin_poseidon_hash_many_cuda(
+  const PoseidonInst* poseidon,
+  const scalar_t* inputs,
+  scalar_t* output,
+  unsigned int number_of_states,
+  unsigned int input_block_len,
+  unsigned int output_len,
+  SpongeConfig* cfg);
+
+cudaError_t grumpkin_poseidon_delete_cuda(PoseidonInst* poseidon);
 
 #ifdef __cplusplus
 }

wrappers/golang/curves/grumpkin/poseidon/poseidon.go

@@ -3,55 +3,85 @@ package poseidon
 // #cgo CFLAGS: -I./include/
 // #include "poseidon.h"
 import "C"
-
 import (
+	"runtime"
 	"unsafe"
 
 	"github.com/ingonyama-zk/icicle/v2/wrappers/golang/core"
 	cr "github.com/ingonyama-zk/icicle/v2/wrappers/golang/cuda_runtime"
+	grumpkin "github.com/ingonyama-zk/icicle/v2/wrappers/golang/curves/grumpkin"
 )
 
-func GetDefaultPoseidonConfig() core.PoseidonConfig {
-	return core.GetDefaultPoseidonConfig()
+type PoseidonHandler = C.struct_PoseidonInst
+type Poseidon struct {
+	width  uint32
+	handle *PoseidonHandler
 }
 
-func PoseidonHash[T any](scalars, results core.HostOrDeviceSlice, numberOfStates int, cfg *core.PoseidonConfig, constants *core.PoseidonConstants[T]) core.IcicleError {
-	scalarsPointer, resultsPointer, cfgPointer := core.PoseidonCheck(scalars, results, cfg, constants, numberOfStates)
+func Create(arity uint32, alpha uint32, fullRoundsHalf uint32, partialRounds uint32, scalars core.HostOrDeviceSlice, mdsMatrix core.HostOrDeviceSlice, nonSparseMatrix core.HostOrDeviceSlice, sparseMatrices core.HostOrDeviceSlice, domainTag grumpkin.ScalarField, ctx *cr.DeviceContext) (*Poseidon, core.IcicleError) {
+	var poseidon *PoseidonHandler
+	cArity := (C.uint)(arity)
+	cAlpha := (C.uint)(alpha)
+	cFullRoundsHalf := (C.uint)(fullRoundsHalf)
+	cPartialRounds := (C.uint)(partialRounds)
+	cScalars := (*C.scalar_t)(scalars.AsUnsafePointer())
+	cMdsMatrix := (*C.scalar_t)(mdsMatrix.AsUnsafePointer())
+	cNonSparseMatrix := (*C.scalar_t)(nonSparseMatrix.AsUnsafePointer())
+	cSparseMatrices := (*C.scalar_t)(sparseMatrices.AsUnsafePointer())
+	cDomainTag := (*C.scalar_t)(unsafe.Pointer(&domainTag))
+	cCtx := (*C.DeviceContext)(unsafe.Pointer(ctx))
+	__ret := C.grumpkin_poseidon_create_cuda(&poseidon, cArity, cAlpha, cFullRoundsHalf, cPartialRounds, cScalars, cMdsMatrix, cNonSparseMatrix, cSparseMatrices, cDomainTag, cCtx)
+	err := core.FromCudaError((cr.CudaError)(__ret))
+	if err.IcicleErrorCode != core.IcicleSuccess {
+		return nil, err
+	}
+	p := Poseidon{handle: poseidon, width: arity + 1}
+	runtime.SetFinalizer(&p, func(p *Poseidon) {
+		p.Delete()
+	})
+	return &p, err
+}
 
-	cScalars := (*C.scalar_t)(scalarsPointer)
-	cResults := (*C.scalar_t)(resultsPointer)
-	cNumberOfStates := (C.int)(numberOfStates)
-	cArity := (C.int)(constants.Arity)
-	cConstants := (*C.PoseidonConstants)(unsafe.Pointer(constants))
-	cCfg := (*C.PoseidonConfig)(cfgPointer)
+func Load(arity uint32, ctx *cr.DeviceContext) (*Poseidon, core.IcicleError) {
+	var poseidon *PoseidonHandler
+	cArity := (C.uint)(arity)
+	cCtx := (*C.DeviceContext)(unsafe.Pointer(ctx))
+	__ret := C.grumpkin_poseidon_load_cuda(&poseidon, cArity, cCtx)
+	err := core.FromCudaError((cr.CudaError)(__ret))
+	if err.IcicleErrorCode != core.IcicleSuccess {
+		return nil, err
+	}
+	p := Poseidon{handle: poseidon, width: arity + 1}
+	runtime.SetFinalizer(&p, func(p *Poseidon) {
+		p.Delete()
+	})
+	return &p, err
+}
 
-	__ret := C.grumpkin_poseidon_hash_cuda(cScalars, cResults, cNumberOfStates, cArity, cConstants, cCfg)
+func (poseidon *Poseidon) HashMany(inputs core.HostOrDeviceSlice, output core.HostOrDeviceSlice, numberOfStates uint32, inputBlockLen uint32, outputLen uint32, cfg *core.SpongeConfig) core.IcicleError {
+	core.SpongeInputCheck(inputs, numberOfStates, inputBlockLen, cfg.InputRate, &cfg.Ctx)
+	core.SpongeOutputsCheck(output, numberOfStates, outputLen, poseidon.width, false, &cfg.Ctx)
 
+	cInputs := (*C.scalar_t)(inputs.AsUnsafePointer())
+	cOutput := (*C.scalar_t)(output.AsUnsafePointer())
+	cNumberOfStates := (C.uint)(numberOfStates)
+	cInputBlockLen := (C.uint)(inputBlockLen)
+	cOutputLen := (C.uint)(outputLen)
+	cCfg := (*C.SpongeConfig)(unsafe.Pointer(cfg))
+	__ret := C.grumpkin_poseidon_hash_many_cuda(poseidon.handle, cInputs, cOutput, cNumberOfStates, cInputBlockLen, cOutputLen, cCfg)
 	err := (cr.CudaError)(__ret)
 	return core.FromCudaError(err)
 }
 
-func CreateOptimizedPoseidonConstants[T any](arity, fullRoundsHalfs, partialRounds int, constants core.HostOrDeviceSlice, ctx cr.DeviceContext, poseidonConstants *core.PoseidonConstants[T]) core.IcicleError {
-
-	cArity := (C.int)(arity)
-	cFullRoundsHalfs := (C.int)(fullRoundsHalfs)
-	cPartialRounds := (C.int)(partialRounds)
-	cConstants := (*C.scalar_t)(constants.AsUnsafePointer())
-	cCtx := (*C.DeviceContext)(unsafe.Pointer(&ctx))
-	cPoseidonConstants := (*C.PoseidonConstants)(unsafe.Pointer(poseidonConstants))
-
-	__ret := C.grumpkin_create_optimized_poseidon_constants_cuda(cArity, cFullRoundsHalfs, cPartialRounds, cConstants, cCtx, cPoseidonConstants)
+func (poseidon *Poseidon) Delete() core.IcicleError {
+	__ret := C.grumpkin_poseidon_delete_cuda(poseidon.handle)
 	err := (cr.CudaError)(__ret)
 	return core.FromCudaError(err)
 }
 
-func InitOptimizedPoseidonConstantsCuda[T any](arity int, ctx cr.DeviceContext, constants *core.PoseidonConstants[T]) core.IcicleError {
-
-	cArity := (C.int)(arity)
-	cCtx := (*C.DeviceContext)(unsafe.Pointer(&ctx))
-	cConstants := (*C.PoseidonConstants)(unsafe.Pointer(constants))
-
-	__ret := C.grumpkin_init_optimized_poseidon_constants_cuda(cArity, cCtx, cConstants)
-	err := (cr.CudaError)(__ret)
-	return core.FromCudaError(err)
+func (poseidon *Poseidon) GetDefaultSpongeConfig() core.SpongeConfig {
+	cfg := core.GetDefaultSpongeConfig()
+	cfg.InputRate = poseidon.width - 1
+	cfg.OutputRate = poseidon.width
+	return cfg
 }

wrappers/golang/curves/grumpkin/tests/poseidon_test.go

@@ -7,6 +7,7 @@ import (
 	cr "github.com/ingonyama-zk/icicle/v2/wrappers/golang/cuda_runtime"
 	grumpkin "github.com/ingonyama-zk/icicle/v2/wrappers/golang/curves/grumpkin"
 	poseidon "github.com/ingonyama-zk/icicle/v2/wrappers/golang/curves/grumpkin/poseidon"
+	"github.com/stretchr/testify/assert"
 )
 
 func TestPoseidon(t *testing.T) {
@@ -14,14 +15,11 @@ func TestPoseidon(t *testing.T) {
 	arity := 2
 	numberOfStates := 1
 
-	cfg := poseidon.GetDefaultPoseidonConfig()
-	cfg.IsAsync = true
-	stream, _ := cr.CreateStream()
-	cfg.Ctx.Stream = &stream
+	ctx, _ := cr.GetDefaultDeviceContext()
+	p, err := poseidon.Load(uint32(arity), &ctx)
+	assert.Equal(t, core.IcicleSuccess, err.IcicleErrorCode)
 
-	var constants core.PoseidonConstants[grumpkin.ScalarField]
-
-	poseidon.InitOptimizedPoseidonConstantsCuda(arity, cfg.Ctx, &constants) //generate constants
+	cfg := p.GetDefaultSpongeConfig()
 
 	scalars := grumpkin.GenerateScalars(numberOfStates * arity)
 	scalars[0] = scalars[0].Zero()
@@ -30,13 +28,13 @@ func TestPoseidon(t *testing.T) {
 	scalarsCopy := core.HostSliceFromElements(scalars[:numberOfStates*arity])
 
 	var deviceInput core.DeviceSlice
-	scalarsCopy.CopyToDeviceAsync(&deviceInput, stream, true)
+	scalarsCopy.CopyToDevice(&deviceInput, true)
 	var deviceOutput core.DeviceSlice
-	deviceOutput.MallocAsync(numberOfStates*scalarsCopy.SizeOfElement(), scalarsCopy.SizeOfElement(), stream)
+	deviceOutput.Malloc(numberOfStates*scalarsCopy.SizeOfElement(), scalarsCopy.SizeOfElement())
 
-	poseidon.PoseidonHash(deviceInput, deviceOutput, numberOfStates, &cfg, &constants) //run Hash function
+	err = p.HashMany(deviceInput, deviceOutput, uint32(numberOfStates), 1, 1, &cfg) //run Hash function
+	assert.Equal(t, core.IcicleSuccess, err.IcicleErrorCode)
 
 	output := make(core.HostSlice[grumpkin.ScalarField], numberOfStates)
-	output.CopyFromDeviceAsync(&deviceOutput, stream)
-
+	output.CopyFromDevice(&deviceOutput)
 }

wrappers/golang/internal/generator/poseidon/templates/poseidon.go.tmpl

@@ -3,55 +3,85 @@ package {{.PackageName}}
 // #cgo CFLAGS: -I./include/
 // #include "poseidon.h"
 import "C"
-
 import (
+	"runtime"
 	"unsafe"
 
 	"github.com/ingonyama-zk/icicle/v2/wrappers/golang/core"
 	cr "github.com/ingonyama-zk/icicle/v2/wrappers/golang/cuda_runtime"
+	{{.Field}} "github.com/ingonyama-zk/icicle/v2/wrappers/golang/{{.BaseImportPath}}"
 )
 
-func GetDefaultPoseidonConfig() core.PoseidonConfig {
-	return core.GetDefaultPoseidonConfig()
+type PoseidonHandler = C.struct_PoseidonInst
+type Poseidon struct {
+	width  uint32
+	handle *PoseidonHandler
 }
 
-func PoseidonHash[T any](scalars, results core.HostOrDeviceSlice, numberOfStates int, cfg *core.PoseidonConfig, constants *core.PoseidonConstants[T]) core.IcicleError {
-	scalarsPointer, resultsPointer, cfgPointer := core.PoseidonCheck(scalars, results, cfg, constants, numberOfStates)
+func Create(arity uint32, alpha uint32, fullRoundsHalf uint32, partialRounds uint32, scalars core.HostOrDeviceSlice, mdsMatrix core.HostOrDeviceSlice, nonSparseMatrix core.HostOrDeviceSlice, sparseMatrices core.HostOrDeviceSlice, domainTag {{.Field}}.ScalarField, ctx *cr.DeviceContext) (*Poseidon, core.IcicleError) {
+	var poseidon *PoseidonHandler
+	cArity := (C.uint)(arity)
+	cAlpha := (C.uint)(alpha)
+	cFullRoundsHalf := (C.uint)(fullRoundsHalf)
+	cPartialRounds := (C.uint)(partialRounds)
+	cScalars := (*C.scalar_t)(scalars.AsUnsafePointer())
+	cMdsMatrix := (*C.scalar_t)(mdsMatrix.AsUnsafePointer())
+	cNonSparseMatrix := (*C.scalar_t)(nonSparseMatrix.AsUnsafePointer())
+	cSparseMatrices := (*C.scalar_t)(sparseMatrices.AsUnsafePointer())
+	cDomainTag := (*C.scalar_t)(unsafe.Pointer(&domainTag))
+	cCtx := (*C.DeviceContext)(unsafe.Pointer(ctx))
+	__ret := C.{{.Field}}_poseidon_create_cuda(&poseidon, cArity, cAlpha, cFullRoundsHalf, cPartialRounds, cScalars, cMdsMatrix, cNonSparseMatrix, cSparseMatrices, cDomainTag, cCtx)
+	err := core.FromCudaError((cr.CudaError)(__ret))
+	if err.IcicleErrorCode != core.IcicleSuccess {
+		return nil, err
+	}
+	p := Poseidon{handle: poseidon, width: arity + 1}
+	runtime.SetFinalizer(&p, func(p *Poseidon) {
+		p.Delete()
+	})
+	return &p, err
+}
 
-	cScalars := (*C.scalar_t)(scalarsPointer)
-	cResults := (*C.scalar_t)(resultsPointer)
-	cNumberOfStates := (C.int)(numberOfStates)
-	cArity := (C.int)(constants.Arity)
-	cConstants := (*C.PoseidonConstants)(unsafe.Pointer(constants))
-	cCfg := (*C.PoseidonConfig)(cfgPointer)
+func Load(arity uint32, ctx *cr.DeviceContext) (*Poseidon, core.IcicleError) {
+	var poseidon *PoseidonHandler
+	cArity := (C.uint)(arity)
+	cCtx := (*C.DeviceContext)(unsafe.Pointer(ctx))
+	__ret := C.{{.Field}}_poseidon_load_cuda(&poseidon, cArity, cCtx)
+	err := core.FromCudaError((cr.CudaError)(__ret))
+	if err.IcicleErrorCode != core.IcicleSuccess {
+		return nil, err
+	}
+	p := Poseidon{handle: poseidon, width: arity + 1}
+	runtime.SetFinalizer(&p, func(p *Poseidon) {
+		p.Delete()
+	})
+	return &p, err
+}
 
-	__ret := C.{{.Field}}_poseidon_hash_cuda(cScalars, cResults, cNumberOfStates, cArity, cConstants, cCfg)
+func (poseidon *Poseidon) HashMany(inputs core.HostOrDeviceSlice, output core.HostOrDeviceSlice, numberOfStates uint32, inputBlockLen uint32, outputLen uint32, cfg *core.SpongeConfig) core.IcicleError {
+	core.SpongeInputCheck(inputs, numberOfStates, inputBlockLen, cfg.InputRate, &cfg.Ctx)
+	core.SpongeOutputsCheck(output, numberOfStates, outputLen, poseidon.width, false, &cfg.Ctx)
 
+	cInputs := (*C.scalar_t)(inputs.AsUnsafePointer())
+	cOutput := (*C.scalar_t)(output.AsUnsafePointer())
+	cNumberOfStates := (C.uint)(numberOfStates)
+	cInputBlockLen := (C.uint)(inputBlockLen)
+	cOutputLen := (C.uint)(outputLen)
+	cCfg := (*C.SpongeConfig)(unsafe.Pointer(cfg))
+	__ret := C.{{.Field}}_poseidon_hash_many_cuda(poseidon.handle, cInputs, cOutput, cNumberOfStates, cInputBlockLen, cOutputLen, cCfg)
 	err := (cr.CudaError)(__ret)
 	return core.FromCudaError(err)
 }
 
-func CreateOptimizedPoseidonConstants[T any](arity, fullRoundsHalfs, partialRounds int, constants core.HostOrDeviceSlice, ctx cr.DeviceContext, poseidonConstants *core.PoseidonConstants[T]) core.IcicleError {
-
-	cArity := (C.int)(arity)
-	cFullRoundsHalfs := (C.int)(fullRoundsHalfs)
-	cPartialRounds := (C.int)(partialRounds)
-	cConstants := (*C.scalar_t)(constants.AsUnsafePointer())
-	cCtx := (*C.DeviceContext)(unsafe.Pointer(&ctx))
-	cPoseidonConstants := (*C.PoseidonConstants)(unsafe.Pointer(poseidonConstants))
-
-	__ret := C.{{.Field}}_create_optimized_poseidon_constants_cuda(cArity, cFullRoundsHalfs, cPartialRounds, cConstants, cCtx, cPoseidonConstants)
+func (poseidon *Poseidon) Delete() core.IcicleError {
+	__ret := C.{{.Field}}_poseidon_delete_cuda(poseidon.handle)
 	err := (cr.CudaError)(__ret)
 	return core.FromCudaError(err)
 }
 
-func InitOptimizedPoseidonConstantsCuda[T any](arity int, ctx cr.DeviceContext, constants *core.PoseidonConstants[T]) core.IcicleError {
-
-	cArity := (C.int)(arity)
-	cCtx := (*C.DeviceContext)(unsafe.Pointer(&ctx))
-	cConstants := (*C.PoseidonConstants)(unsafe.Pointer(constants))
-
-	__ret := C.{{.Field}}_init_optimized_poseidon_constants_cuda(cArity, cCtx, cConstants)
-	err := (cr.CudaError)(__ret)
-	return core.FromCudaError(err)
+func (poseidon *Poseidon) GetDefaultSpongeConfig() core.SpongeConfig {
+	cfg := core.GetDefaultSpongeConfig()
+	cfg.InputRate = poseidon.width - 1
+	cfg.OutputRate = poseidon.width
+	return cfg
 }

wrappers/golang/internal/generator/poseidon/templates/poseidon.h.tmpl

@@ -9,14 +9,40 @@ extern "C" {
 #endif
 
 typedef struct scalar_t scalar_t;
-typedef struct PoseidonConfig PoseidonConfig;
 typedef struct DeviceContext DeviceContext;
-typedef struct PoseidonConstants PoseidonConstants;
+typedef struct TreeBuilderConfig TreeBuilderConfig;
+typedef struct PoseidonInst PoseidonInst;
+typedef struct SpongeConfig SpongeConfig;
 
 
-cudaError_t {{.Field}}_poseidon_hash_cuda(const scalar_t* input, scalar_t* output, int number_of_states, int arity, PoseidonConstants* constants, PoseidonConfig* config);
-cudaError_t {{.Field}}_create_optimized_poseidon_constants_cuda(int arity, int full_rounds_halfs, int partial_rounds, const scalar_t* constants, DeviceContext* ctx, PoseidonConstants* poseidon_constants);
-cudaError_t {{.Field}}_init_optimized_poseidon_constants_cuda(int arity, DeviceContext* ctx, PoseidonConstants* constants);
+cudaError_t {{.Field}}_poseidon_create_cuda(
+  PoseidonInst** poseidon,
+  unsigned int arity,
+  unsigned int alpha,
+  unsigned int partial_rounds,
+  unsigned int full_rounds_half,
+  const scalar_t* round_constants,
+  const scalar_t* mds_matrix,
+  const scalar_t* non_sparse_matrix,
+  const scalar_t* sparse_matrices,
+  const scalar_t* domain_tag,
+  DeviceContext* ctx);
+
+cudaError_t {{.Field}}_poseidon_load_cuda(
+  PoseidonInst** poseidon,
+  unsigned int arity,
+  DeviceContext* ctx);
+
+cudaError_t {{.Field}}_poseidon_hash_many_cuda(
+  const PoseidonInst* poseidon,
+  const scalar_t* inputs,
+  scalar_t* output,
+  unsigned int number_of_states,
+  unsigned int input_block_len,
+  unsigned int output_len,
+  SpongeConfig* cfg);
+
+cudaError_t {{.Field}}_poseidon_delete_cuda(PoseidonInst* poseidon);
 
 #ifdef __cplusplus
 }

wrappers/golang/internal/generator/poseidon/templates/poseidon_test.go.tmpl

@@ -2,37 +2,24 @@ package tests
 
 import (
 	"testing"
-	
+
 	core "github.com/ingonyama-zk/icicle/v2/wrappers/golang/core"
 	cr "github.com/ingonyama-zk/icicle/v2/wrappers/golang/cuda_runtime"
 	{{.Field}} "github.com/ingonyama-zk/icicle/v2/wrappers/golang/{{.BaseImportPath}}"
 	poseidon "github.com/ingonyama-zk/icicle/v2/wrappers/golang/{{.BaseImportPath}}/poseidon"
-
-	{{if eq .Field "bls12_381"}}
-	"fmt"
 	"github.com/stretchr/testify/assert"
-	{{end}}
 )
-{{if eq .Field "bls12_381"}}
-func formatOutput(x {{.Field}}.{{.FieldPrefix}}Field) string {
-	r := x.GetLimbs()
-	return fmt.Sprintf("%08x%08x%08x%08x%08x%08x%08x%08x", r[7], r[6], r[5], r[4], r[3], r[2], r[1], r[0])
-}
-{{end}}
 
 func TestPoseidon(t *testing.T) {
 
 	arity := 2
 	numberOfStates := 1
 
-	cfg := poseidon.GetDefaultPoseidonConfig()
-	cfg.IsAsync = true
-	stream, _ := cr.CreateStream()
-	cfg.Ctx.Stream = &stream
+	ctx, _ := cr.GetDefaultDeviceContext()
+	p, err := poseidon.Load(uint32(arity), &ctx)
+	assert.Equal(t, core.IcicleSuccess, err.IcicleErrorCode)
 
-	var constants core.PoseidonConstants[{{.Field}}.{{.FieldPrefix}}Field]
-
-	poseidon.InitOptimizedPoseidonConstantsCuda(arity, cfg.Ctx, &constants) //generate constants
+	cfg := p.GetDefaultSpongeConfig()
 
 	scalars := {{.Field}}.GenerateScalars(numberOfStates * arity)
 	scalars[0] = scalars[0].Zero()
@@ -41,19 +28,13 @@ func TestPoseidon(t *testing.T) {
 	scalarsCopy := core.HostSliceFromElements(scalars[:numberOfStates*arity])
 
 	var deviceInput core.DeviceSlice
-	scalarsCopy.CopyToDeviceAsync(&deviceInput, stream, true)
+	scalarsCopy.CopyToDevice(&deviceInput, true)
 	var deviceOutput core.DeviceSlice
-	deviceOutput.MallocAsync(numberOfStates*scalarsCopy.SizeOfElement(), scalarsCopy.SizeOfElement(), stream)
+	deviceOutput.Malloc(numberOfStates*scalarsCopy.SizeOfElement(), scalarsCopy.SizeOfElement())
 
-	poseidon.PoseidonHash(deviceInput, deviceOutput, numberOfStates, &cfg, &constants) //run Hash function
+	err = p.HashMany(deviceInput, deviceOutput, uint32(numberOfStates), 1, 1, &cfg) //run Hash function
+	assert.Equal(t, core.IcicleSuccess, err.IcicleErrorCode)
 
-	output := make(core.HostSlice[{{.Field}}.{{.FieldPrefix}}Field], numberOfStates)
-	output.CopyFromDeviceAsync(&deviceOutput, stream)
-
-	{{if eq .Field "bls12_381"}}
-	expectedString := "48fe0b1331196f6cdb33a7c6e5af61b76fd388e1ef1d3d418be5147f0e4613d4" //This result is from https://github.com/triplewz/poseidon
-	outputString := formatOutput(output[0])
-
-	assert.Equal(t, outputString, expectedString, "Poseidon hash does not match expected result")
-	{{end}}
+	output := make(core.HostSlice[{{.Field}}.ScalarField], numberOfStates)
+	output.CopyFromDevice(&deviceOutput)
 }

wrappers/rust/icicle-core/src/hash.rs

@@ -0,0 +1,136 @@
+use std::ffi::c_void;
+
+use icicle_cuda_runtime::{
+    device::check_device,
+    device_context::{DeviceContext, DEFAULT_DEVICE_ID},
+    memory::HostOrDeviceSlice,
+};
+
+use crate::ntt::IcicleResult;
+
+/// Struct that encodes Sponge hash parameters.
+#[repr(C)]
+#[derive(Debug, Clone)]
+pub struct SpongeConfig<'a> {
+    /// Details related to the device such as its id and stream id. See [DeviceContext](@ref device_context::DeviceContext).
+    pub ctx: DeviceContext<'a>,
+    pub(crate) are_inputs_on_device: bool,
+    pub(crate) are_outputs_on_device: bool,
+    pub input_rate: u32,
+    pub output_rate: u32,
+    pub offset: u32,
+
+    /// If true - input should be already aligned for poseidon permutation.
+    /// Aligned format: [0, A, B, 0, C, D, ...] (as you might get by using loop_state)
+    /// not aligned format: [A, B, 0, C, D, 0, ...] (as you might get from cudaMemcpy2D)
+    pub recursive_squeeze: bool,
+
+    /// If true, hash results will also be copied in the input pointer in aligned format
+    pub aligned: bool,
+    /// Whether to run the sponge operations asynchronously. If set to `true`, the functions will be non-blocking and you'd need to synchronize
+    /// it explicitly by running `stream.synchronize()`. If set to false, the functions will block the current CPU thread.
+    pub is_async: bool,
+}
+
+impl<'a> Default for SpongeConfig<'a> {
+    fn default() -> Self {
+        Self::default_for_device(DEFAULT_DEVICE_ID)
+    }
+}
+
+impl<'a> SpongeConfig<'a> {
+    pub(crate) fn default_for_device(device_id: usize) -> Self {
+        SpongeConfig {
+            ctx: DeviceContext::default_for_device(device_id),
+            are_inputs_on_device: false,
+            are_outputs_on_device: false,
+            input_rate: 0,
+            output_rate: 0,
+            offset: 0,
+            recursive_squeeze: false,
+            aligned: false,
+            is_async: false,
+        }
+    }
+}
+
+pub trait SpongeHash<PreImage, Image> {
+    fn hash_many(
+        &self,
+        inputs: &(impl HostOrDeviceSlice<PreImage> + ?Sized),
+        output: &mut (impl HostOrDeviceSlice<Image> + ?Sized),
+        number_of_states: usize,
+        input_block_len: usize,
+        output_len: usize,
+        cfg: &SpongeConfig,
+    ) -> IcicleResult<()>;
+
+    fn default_config<'a>(&self) -> SpongeConfig<'a>;
+
+    fn get_handle(&self) -> *const c_void;
+}
+
+pub(crate) fn sponge_check_input<T>(
+    inputs: &(impl HostOrDeviceSlice<T> + ?Sized),
+    number_of_states: usize,
+    input_block_len: usize,
+    input_rate: usize,
+    ctx: &DeviceContext,
+) {
+    if input_block_len > input_rate {
+        panic!(
+            "input block len ({}) can't be greater than input rate ({})",
+            input_block_len, input_rate
+        );
+    }
+
+    let inputs_size_expected = input_block_len * number_of_states;
+    if inputs.len() < inputs_size_expected {
+        panic!(
+            "inputs len is {}; but needs to be at least {}",
+            inputs.len(),
+            inputs_size_expected,
+        );
+    }
+
+    let ctx_device_id = ctx.device_id;
+    if let Some(device_id) = inputs.device_id() {
+        assert_eq!(
+            device_id, ctx_device_id,
+            "Device ids in inputs and context are different"
+        );
+    }
+    check_device(ctx_device_id);
+}
+
+pub(crate) fn sponge_check_outputs<T>(
+    outputs: &(impl HostOrDeviceSlice<T> + ?Sized),
+    number_of_states: usize,
+    output_len: usize,
+    width: usize,
+    recursive: bool,
+    ctx: &DeviceContext,
+) {
+    let outputs_size_expected = if recursive {
+        width * number_of_states
+    } else {
+        output_len * number_of_states
+    };
+
+    if outputs.len() < outputs_size_expected {
+        panic!(
+            "outputs len is {}; but needs to be at least {}",
+            outputs.len(),
+            outputs_size_expected,
+        );
+    }
+
+    let ctx_device_id = ctx.device_id;
+    if let Some(device_id) = outputs.device_id() {
+        assert_eq!(
+            device_id, ctx_device_id,
+            "Device ids in outputs and context are different"
+        );
+    }
+    check_device(ctx_device_id);
+}

wrappers/rust/icicle-core/src/lib.rs

@@ -1,7 +1,10 @@
+use std::ffi::c_void;
+
 pub mod curve;
 pub mod ecntt;
 pub mod error;
 pub mod field;
+pub mod hash;
 pub mod msm;
 pub mod ntt;
 pub mod polynomials;
@@ -18,3 +21,11 @@ where
     <Self::ScalarField as traits::FieldImpl>::Config: ntt::NTT<Self::ScalarField, Self::ScalarField>,
 {
 }
+
+#[repr(C)]
+#[derive(Debug)]
+pub struct Matrix {
+    pub values: *const c_void,
+    pub width: usize,
+    pub height: usize,
+}

wrappers/rust/icicle-core/src/poseidon/mod.rs

@@ -1,212 +1,157 @@
 #[doc(hidden)]
 pub mod tests;
 
-use icicle_cuda_runtime::{
-    device::check_device,
-    device_context::{DeviceContext, DEFAULT_DEVICE_ID},
-    memory::{DeviceSlice, HostOrDeviceSlice},
+use std::{ffi::c_void, marker::PhantomData};
+
+use icicle_cuda_runtime::{device_context::DeviceContext, memory::HostOrDeviceSlice};
+
+use crate::{
+    error::IcicleResult,
+    hash::{sponge_check_input, sponge_check_outputs, SpongeConfig, SpongeHash},
+    traits::FieldImpl,
 };
 
-use crate::{error::IcicleResult, traits::FieldImpl};
-
-#[repr(C)]
-pub struct PoseidonConstants<'a, F: FieldImpl> {
-    arity: u32,
-
-    partial_rounds: u32,
-
-    full_rounds_half: u32,
-
-    /// These should be pointers to data allocated on device
-    round_constants: &'a DeviceSlice<F>,
-    mds_matrix: &'a DeviceSlice<F>,
-    non_sparse_matrix: &'a DeviceSlice<F>,
-    sparse_matrices: &'a DeviceSlice<F>,
-
-    /// Domain tag is the first element in the Poseidon state.
-    /// For the Merkle tree mode it should equal 2^arity - 1
-    domain_tag: F,
+pub type PoseidonHandle = *const c_void;
+pub struct Poseidon<F>
+where
+    F: FieldImpl,
+    <F as FieldImpl>::Config: PoseidonImpl<F>,
+{
+    width: usize,
+    handle: PoseidonHandle,
+    phantom: PhantomData<F>,
 }
 
-/// Struct that encodes Poseidon parameters to be passed into the [poseidon_hash_many](poseidon_hash_many) function.
-#[repr(C)]
-#[derive(Debug, Clone)]
-pub struct PoseidonConfig<'a> {
-    /// Details related to the device such as its id and stream id. See [DeviceContext](@ref device_context::DeviceContext).
-    pub ctx: DeviceContext<'a>,
-
-    are_inputs_on_device: bool,
-
-    are_outputs_on_device: bool,
-
-    /// If true, input is considered to be a states vector, holding the preimages
-    /// in aligned or not aligned format. Memory under the input pointer will be used for states
-    /// If false, fresh states memory will be allocated and input will be copied into it
-    pub input_is_a_state: bool,
-
-    /// If true - input should be already aligned for poseidon permutation.
-    /// Aligned format: [0, A, B, 0, C, D, ...] (as you might get by using loop_state)
-    /// not aligned format: [A, B, 0, C, D, 0, ...] (as you might get from cudaMemcpy2D)
-    pub aligned: bool,
-
-    /// If true, hash results will also be copied in the input pointer in aligned format
-    pub loop_state: bool,
-
-    /// Whether to run Poseidon asynchronously. If set to `true`, Poseidon will be non-blocking
-    /// and you'd need to synchronize it explicitly by running `cudaStreamSynchronize` or `cudaDeviceSynchronize`.
-    /// If set to `false`, Poseidon will block the current CPU thread.
-    pub is_async: bool,
-}
-
-impl<'a> Default for PoseidonConfig<'a> {
-    fn default() -> Self {
-        Self::default_for_device(DEFAULT_DEVICE_ID)
+impl<F> Poseidon<F>
+where
+    F: FieldImpl,
+    <F as FieldImpl>::Config: PoseidonImpl<F>,
+{
+    pub fn load(arity: usize, ctx: &DeviceContext) -> IcicleResult<Self> {
+        <<F as FieldImpl>::Config as PoseidonImpl<F>>::load(arity as u32, ctx).and_then(|handle| {
+            Ok(Self {
+                width: arity + 1,
+                handle,
+                phantom: PhantomData,
+            })
+        })
     }
-}
 
-impl<'a> PoseidonConfig<'a> {
-    pub fn default_for_device(device_id: usize) -> Self {
-        Self {
-            ctx: DeviceContext::default_for_device(device_id),
-            are_inputs_on_device: false,
-            are_outputs_on_device: false,
-            input_is_a_state: false,
-            aligned: false,
-            loop_state: false,
-            is_async: false,
-        }
-    }
-}
-
-pub trait Poseidon<F: FieldImpl> {
-    fn create_optimized_constants<'a>(
-        arity: u32,
+    pub fn new(
+        arity: usize,
+        alpha: u32,
         full_rounds_half: u32,
         partial_rounds: u32,
-        constants: &mut [F],
+        round_constants: &[F],
+        mds_matrix: &[F],
+        non_sparse_matrix: &[F],
+        sparse_matrices: &[F],
+        domain_tag: F,
         ctx: &DeviceContext,
-    ) -> IcicleResult<PoseidonConstants<'a, F>>;
-    fn load_optimized_constants<'a>(arity: u32, ctx: &DeviceContext) -> IcicleResult<PoseidonConstants<'a, F>>;
-    fn poseidon_unchecked(
-        input: &mut (impl HostOrDeviceSlice<F> + ?Sized),
+    ) -> IcicleResult<Self> {
+        <<F as FieldImpl>::Config as PoseidonImpl<F>>::create(
+            arity as u32,
+            alpha,
+            full_rounds_half,
+            partial_rounds,
+            round_constants,
+            mds_matrix,
+            non_sparse_matrix,
+            sparse_matrices,
+            domain_tag,
+            ctx,
+        )
+        .and_then(|handle| {
+            Ok(Self {
+                width: arity + 1,
+                handle,
+                phantom: PhantomData,
+            })
+        })
+    }
+}
+
+impl<F> SpongeHash<F, F> for Poseidon<F>
+where
+    F: FieldImpl,
+    <F as FieldImpl>::Config: PoseidonImpl<F>,
+{
+    fn get_handle(&self) -> *const c_void {
+        self.handle
+    }
+
+    fn hash_many(
+        &self,
+        inputs: &(impl HostOrDeviceSlice<F> + ?Sized),
+        output: &mut (impl HostOrDeviceSlice<F> + ?Sized),
+        number_of_states: usize,
+        input_block_len: usize,
+        output_len: usize,
+        cfg: &SpongeConfig,
+    ) -> IcicleResult<()> {
+        sponge_check_input(inputs, number_of_states, input_block_len, self.width - 1, &cfg.ctx);
+        sponge_check_outputs(output, number_of_states, output_len, self.width, false, &cfg.ctx);
+
+        let mut local_cfg = cfg.clone();
+        local_cfg.are_inputs_on_device = inputs.is_on_device();
+        local_cfg.are_outputs_on_device = output.is_on_device();
+
+        <<F as FieldImpl>::Config as PoseidonImpl<F>>::hash_many(
+            inputs,
+            output,
+            number_of_states as u32,
+            input_block_len as u32,
+            output_len as u32,
+            self.handle,
+            &local_cfg,
+        )
+    }
+
+    fn default_config<'a>(&self) -> SpongeConfig<'a> {
+        let mut cfg = SpongeConfig::default();
+        cfg.input_rate = self.width as u32 - 1;
+        cfg.output_rate = self.width as u32;
+        cfg
+    }
+}
+
+impl<F> Drop for Poseidon<F>
+where
+    F: FieldImpl,
+    <F as FieldImpl>::Config: PoseidonImpl<F>,
+{
+    fn drop(&mut self) {
+        <<F as FieldImpl>::Config as PoseidonImpl<F>>::delete(self.handle).unwrap();
+    }
+}
+
+pub trait PoseidonImpl<F: FieldImpl> {
+    fn create(
+        arity: u32,
+        alpha: u32,
+        full_rounds_half: u32,
+        partial_rounds: u32,
+        round_constants: &[F],
+        mds_matrix: &[F],
+        non_sparse_matrix: &[F],
+        sparse_matrices: &[F],
+        domain_tag: F,
+        ctx: &DeviceContext,
+    ) -> IcicleResult<PoseidonHandle>;
+
+    fn load(arity: u32, ctx: &DeviceContext) -> IcicleResult<PoseidonHandle>;
+
+    fn hash_many(
+        inputs: &(impl HostOrDeviceSlice<F> + ?Sized),
         output: &mut (impl HostOrDeviceSlice<F> + ?Sized),
         number_of_states: u32,
-        arity: u32,
-        constants: &PoseidonConstants<F>,
-        config: &PoseidonConfig,
+        input_block_len: u32,
+        output_len: u32,
+        poseidon: PoseidonHandle,
+        cfg: &SpongeConfig,
     ) -> IcicleResult<()>;
-}
 
-/// Loads pre-calculated poseidon constants on the GPU.
-pub fn load_optimized_poseidon_constants<'a, F>(
-    arity: u32,
-    ctx: &DeviceContext,
-) -> IcicleResult<PoseidonConstants<'a, F>>
-where
-    F: FieldImpl,
-    <F as FieldImpl>::Config: Poseidon<F>,
-{
-    <<F as FieldImpl>::Config as Poseidon<F>>::load_optimized_constants(arity, ctx)
-}
-
-/// Creates new instance of poseidon constants on the GPU.
-pub fn create_optimized_poseidon_constants<'a, F>(
-    arity: u32,
-    ctx: &DeviceContext,
-    full_rounds_half: u32,
-    partial_rounds: u32,
-    constants: &mut [F],
-) -> IcicleResult<PoseidonConstants<'a, F>>
-where
-    F: FieldImpl,
-    <F as FieldImpl>::Config: Poseidon<F>,
-{
-    <<F as FieldImpl>::Config as Poseidon<F>>::create_optimized_constants(
-        arity,
-        full_rounds_half,
-        partial_rounds,
-        constants,
-        ctx,
-    )
-}
-
-/// Computes the poseidon hashes for multiple preimages.
-///
-/// # Arguments
-///
-/// * `input` - a pointer to the input data. May point to a vector of preimages or a vector of states filled with preimages.
-///
-/// * `output` - a pointer to the output data. Must be at least of size [number_of_states](number_of_states)
-///
-/// * `number_of_states` - number of input blocks of size `arity`
-///
-/// * `arity` - the arity of the hash function (the size of 1 preimage)
-///
-/// * `constants` - Poseidon constants.
-///
-/// * `config` - config used to specify extra arguments of the Poseidon.
-pub fn poseidon_hash_many<F>(
-    input: &mut (impl HostOrDeviceSlice<F> + ?Sized),
-    output: &mut (impl HostOrDeviceSlice<F> + ?Sized),
-    number_of_states: u32,
-    arity: u32,
-    constants: &PoseidonConstants<F>,
-    config: &PoseidonConfig,
-) -> IcicleResult<()>
-where
-    F: FieldImpl,
-    <F as FieldImpl>::Config: Poseidon<F>,
-{
-    let input_len_required = if config.input_is_a_state {
-        number_of_states * (arity + 1)
-    } else {
-        number_of_states * arity
-    };
-
-    if input.len() < input_len_required as usize {
-        panic!(
-            "input len is {}; but needs to be at least {}",
-            input.len(),
-            input_len_required
-        );
-    }
-
-    if output.len() < number_of_states as usize {
-        panic!(
-            "output len is {}; but needs to be at least {}",
-            output.len(),
-            number_of_states
-        );
-    }
-
-    let ctx_device_id = config
-        .ctx
-        .device_id;
-    if let Some(device_id) = input.device_id() {
-        assert_eq!(
-            device_id, ctx_device_id,
-            "Device ids in input and context are different"
-        );
-    }
-    if let Some(device_id) = output.device_id() {
-        assert_eq!(
-            device_id, ctx_device_id,
-            "Device ids in output and context are different"
-        );
-    }
-    check_device(ctx_device_id);
-    let mut local_cfg = config.clone();
-    local_cfg.are_inputs_on_device = input.is_on_device();
-    local_cfg.are_outputs_on_device = output.is_on_device();
-
-    <<F as FieldImpl>::Config as Poseidon<F>>::poseidon_unchecked(
-        input,
-        output,
-        number_of_states,
-        arity,
-        constants,
-        &local_cfg,
-    )
+    fn delete(poseidon: PoseidonHandle) -> IcicleResult<()>;
 }
 
 #[macro_export]
@@ -218,91 +163,110 @@ macro_rules! impl_poseidon {
       $field_config:ident
     ) => {
         mod $field_prefix_ident {
-            use crate::poseidon::{$field, $field_config, CudaError, DeviceContext, PoseidonConfig, PoseidonConstants};
+            use crate::poseidon::{$field, $field_config, CudaError, DeviceContext, PoseidonHandle, SpongeConfig};
             extern "C" {
-                #[link_name = concat!($field_prefix, "_create_optimized_poseidon_constants_cuda")]
-                pub(crate) fn _create_optimized_constants(
+                #[link_name = concat!($field_prefix, "_poseidon_create_cuda")]
+                pub(crate) fn create(
+                    poseidon: *mut PoseidonHandle,
                     arity: u32,
+                    alpha: u32,
                     full_rounds_half: u32,
                     partial_rounds: u32,
-                    constants: *mut $field,
+                    round_constants: *const $field,
+                    mds_matrix: *const $field,
+                    non_sparse_matrix: *const $field,
+                    sparse_matrices: *const $field,
+                    domain_tag: $field,
                     ctx: &DeviceContext,
-                    poseidon_constants: *mut PoseidonConstants<$field>,
                 ) -> CudaError;
 
-                #[link_name = concat!($field_prefix, "_init_optimized_poseidon_constants_cuda")]
-                pub(crate) fn _load_optimized_constants(
-                    arity: u32,
-                    ctx: &DeviceContext,
-                    constants: *mut PoseidonConstants<$field>,
-                ) -> CudaError;
+                #[link_name = concat!($field_prefix, "_poseidon_load_cuda")]
+                pub(crate) fn load(poseidon: *mut PoseidonHandle, arity: u32, ctx: &DeviceContext) -> CudaError;
 
-                #[link_name = concat!($field_prefix, "_poseidon_hash_cuda")]
+                #[link_name = concat!($field_prefix, "_poseidon_delete_cuda")]
+                pub(crate) fn delete(poseidon: PoseidonHandle) -> CudaError;
+
+                #[link_name = concat!($field_prefix, "_poseidon_hash_many_cuda")]
                 pub(crate) fn hash_many(
-                    input: *mut $field,
+                    poseidon: PoseidonHandle,
+                    inputs: *const $field,
                     output: *mut $field,
                     number_of_states: u32,
-                    arity: u32,
-                    constants: &PoseidonConstants<$field>,
-                    config: &PoseidonConfig,
+                    input_block_len: u32,
+                    output_len: u32,
+                    cfg: &SpongeConfig,
                 ) -> CudaError;
             }
         }
 
-        impl Poseidon<$field> for $field_config {
-            fn create_optimized_constants<'a>(
+        impl PoseidonImpl<$field> for $field_config {
+            fn create(
                 arity: u32,
+                alpha: u32,
                 full_rounds_half: u32,
                 partial_rounds: u32,
-                constants: &mut [$field],
+                round_constants: &[$field],
+                mds_matrix: &[$field],
+                non_sparse_matrix: &[$field],
+                sparse_matrices: &[$field],
+                domain_tag: $field,
                 ctx: &DeviceContext,
-            ) -> IcicleResult<PoseidonConstants<'a, $field>> {
+            ) -> IcicleResult<PoseidonHandle> {
                 unsafe {
-                    let mut poseidon_constants = MaybeUninit::<PoseidonConstants<'a, $field>>::uninit();
-                    let err = $field_prefix_ident::_create_optimized_constants(
+                    let mut poseidon = MaybeUninit::<PoseidonHandle>::uninit();
+                    $field_prefix_ident::create(
+                        poseidon.as_mut_ptr(),
                         arity,
+                        alpha,
                         full_rounds_half,
                         partial_rounds,
-                        constants as *mut _ as *mut $field,
+                        round_constants as *const _ as *const $field,
+                        mds_matrix as *const _ as *const $field,
+                        non_sparse_matrix as *const _ as *const $field,
+                        sparse_matrices as *const _ as *const $field,
+                        domain_tag,
                         ctx,
-                        poseidon_constants.as_mut_ptr(),
                     )
-                    .wrap();
-                    err.and(Ok(poseidon_constants.assume_init()))
+                    .wrap()
+                    .and(Ok(poseidon.assume_init()))
                 }
             }
 
-            fn load_optimized_constants<'a>(
-                arity: u32,
-                ctx: &DeviceContext,
-            ) -> IcicleResult<PoseidonConstants<'a, $field>> {
+            fn load(arity: u32, ctx: &DeviceContext) -> IcicleResult<PoseidonHandle> {
                 unsafe {
-                    let mut constants = MaybeUninit::<PoseidonConstants<'a, $field>>::uninit();
-                    let err = $field_prefix_ident::_load_optimized_constants(arity, ctx, constants.as_mut_ptr()).wrap();
-                    err.and(Ok(constants.assume_init()))
+                    let mut poseidon = MaybeUninit::<PoseidonHandle>::uninit();
+                    $field_prefix_ident::load(poseidon.as_mut_ptr(), arity, ctx)
+                        .wrap()
+                        .and(Ok(poseidon.assume_init()))
                 }
             }
 
-            fn poseidon_unchecked(
-                input: &mut (impl HostOrDeviceSlice<$field> + ?Sized),
+            fn hash_many(
+                inputs: &(impl HostOrDeviceSlice<$field> + ?Sized),
                 output: &mut (impl HostOrDeviceSlice<$field> + ?Sized),
                 number_of_states: u32,
-                arity: u32,
-                constants: &PoseidonConstants<$field>,
-                config: &PoseidonConfig,
+                input_block_len: u32,
+                output_len: u32,
+                poseidon: PoseidonHandle,
+                cfg: &SpongeConfig,
             ) -> IcicleResult<()> {
                 unsafe {
                     $field_prefix_ident::hash_many(
-                        input.as_mut_ptr(),
+                        poseidon,
+                        inputs.as_ptr(),
                         output.as_mut_ptr(),
                         number_of_states,
-                        arity,
-                        constants,
-                        config,
+                        input_block_len,
+                        output_len,
+                        cfg,
                     )
                     .wrap()
                 }
             }
+
+            fn delete(poseidon: PoseidonHandle) -> IcicleResult<()> {
+                unsafe { $field_prefix_ident::delete(poseidon).wrap() }
+            }
         }
     };
 }
@@ -318,18 +282,3 @@ macro_rules! impl_poseidon_tests {
         }
     };
 }
-
-#[macro_export]
-macro_rules! impl_poseidon_custom_config_test {
-    (
-      $field:ident,
-      $field_bytes:literal,
-      $field_prefix:literal,
-      $partial_rounds:literal
-    ) => {
-        #[test]
-        fn test_poseidon_custom_config() {
-            check_poseidon_custom_config::<$field>($field_bytes, $field_prefix, $partial_rounds)
-        }
-    };
-}

wrappers/rust/icicle-core/src/poseidon/tests.rs

@@ -1,105 +1,48 @@
+use crate::hash::SpongeHash;
 use crate::traits::FieldImpl;
 use icicle_cuda_runtime::device_context::DeviceContext;
 use icicle_cuda_runtime::memory::{HostOrDeviceSlice, HostSlice};
 
-use std::io::Read;
-use std::path::PathBuf;
-use std::{env, fs::File};
+use super::{Poseidon, PoseidonImpl};
 
-use super::{
-    create_optimized_poseidon_constants, load_optimized_poseidon_constants, poseidon_hash_many, Poseidon,
-    PoseidonConfig, PoseidonConstants,
-};
-
-pub fn init_poseidon<'a, F: FieldImpl>(arity: u32) -> PoseidonConstants<'a, F>
+pub fn init_poseidon<F: FieldImpl>(arity: usize) -> Poseidon<F>
 where
-    <F as FieldImpl>::Config: Poseidon<F>,
+    <F as FieldImpl>::Config: PoseidonImpl<F>,
 {
     let ctx = DeviceContext::default();
-
-    load_optimized_poseidon_constants::<F>(arity, &ctx).unwrap()
+    Poseidon::load(arity, &ctx).unwrap()
 }
 
-pub fn _check_poseidon_hash_many<F: FieldImpl>(constants: PoseidonConstants<F>) -> (F, F)
+pub fn _check_poseidon_hash_many<F: FieldImpl>(poseidon: Poseidon<F>)
 where
-    <F as FieldImpl>::Config: Poseidon<F>,
+    <F as FieldImpl>::Config: PoseidonImpl<F>,
 {
     let test_size = 1 << 10;
-    let arity = 2u32;
-    let mut inputs = vec![F::one(); test_size * arity as usize];
+    let arity = poseidon.width - 1;
+    let mut inputs = vec![F::one(); test_size * arity];
     let mut outputs = vec![F::zero(); test_size];
 
     let input_slice = HostSlice::from_mut_slice(&mut inputs);
     let output_slice = HostSlice::from_mut_slice(&mut outputs);
 
-    let config = PoseidonConfig::default();
-    poseidon_hash_many::<F>(
-        input_slice,
-        output_slice,
-        test_size as u32,
-        arity as u32,
-        &constants,
-        &config,
-    )
-    .unwrap();
+    let cfg = poseidon.default_config();
+    poseidon
+        .hash_many(input_slice, output_slice, test_size, arity, 1, &cfg)
+        .unwrap();
 
     let a1 = output_slice[0];
-    let a2 = output_slice[output_slice.len() - 2];
+    let a2 = output_slice[output_slice.len() - 1];
 
-    println!("first: {:?}, last: {:?}", a1, a2);
     assert_eq!(a1, a2);
-
-    (a1, a2)
 }
 
 pub fn check_poseidon_hash_many<F: FieldImpl>()
 where
-    <F as FieldImpl>::Config: Poseidon<F>,
+    <F as FieldImpl>::Config: PoseidonImpl<F>,
 {
-    for arity in [2, 4] {
-        let constants = init_poseidon::<F>(arity as u32);
+    for arity in [2, 4, 8, 11] {
+        let poseidon = init_poseidon::<F>(arity);
 
-        _check_poseidon_hash_many(constants);
+        _check_poseidon_hash_many(poseidon);
     }
 }
-
-pub fn check_poseidon_custom_config<F: FieldImpl>(field_bytes: usize, field_prefix: &str, partial_rounds: u32)
-where
-    <F as FieldImpl>::Config: Poseidon<F>,
-{
-    let arity = 2u32;
-    let constants = init_poseidon::<F>(arity as u32);
-
-    let full_rounds_half = 4;
-
-    let ctx = DeviceContext::default();
-    let cargo_manifest_dir = env!("CARGO_MANIFEST_DIR");
-    let constants_file = PathBuf::from(cargo_manifest_dir)
-        .join("tests")
-        .join(format!("{}_constants.bin", field_prefix));
-    let mut constants_buf = vec![];
-    File::open(constants_file)
-        .unwrap()
-        .read_to_end(&mut constants_buf)
-        .unwrap();
-
-    let mut custom_constants = vec![];
-    for chunk in constants_buf.chunks(field_bytes) {
-        custom_constants.push(F::from_bytes_le(chunk));
-    }
-
-    let custom_constants = create_optimized_poseidon_constants::<F>(
-        arity as u32,
-        &ctx,
-        full_rounds_half,
-        partial_rounds,
-        &mut custom_constants,
-    )
-    .unwrap();
-
-    let (a1, a2) = _check_poseidon_hash_many(constants);
-    let (b1, b2) = _check_poseidon_hash_many(custom_constants);
-
-    assert_eq!(a1, b1);
-    assert_eq!(a2, b2);
-}

wrappers/rust/icicle-core/src/poseidon2/mod.rs

@@ -1,107 +1,66 @@
 #[doc(hidden)]
 pub mod tests;
 
-use icicle_cuda_runtime::{
-    device::check_device,
-    device_context::{DeviceContext, DEFAULT_DEVICE_ID},
-    memory::{DeviceSlice, HostOrDeviceSlice},
+use std::{ffi::c_void, marker::PhantomData};
+
+use icicle_cuda_runtime::{device_context::DeviceContext, memory::HostOrDeviceSlice};
+
+use crate::{
+    error::IcicleResult,
+    hash::{sponge_check_input, sponge_check_outputs, SpongeConfig, SpongeHash},
+    traits::FieldImpl,
 };
 
-use crate::{error::IcicleResult, traits::FieldImpl};
-
 #[repr(C)]
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, Copy)]
 pub enum DiffusionStrategy {
     Default,
     Montgomery,
 }
 
 #[repr(C)]
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, Copy)]
 pub enum MdsType {
     Default,
     Plonky,
 }
 
-#[repr(C)]
-#[derive(Debug, Clone)]
-pub enum PoseidonMode {
-    Compression,
-    Permutation,
+pub type Poseidon2Handle = *const c_void;
+pub struct Poseidon2<F>
+where
+    F: FieldImpl,
+    <F as FieldImpl>::Config: Poseidon2Impl<F>,
+{
+    width: usize,
+    handle: Poseidon2Handle,
+    phantom: PhantomData<F>,
 }
 
-#[repr(C)]
-pub struct Poseidon2Constants<'a, F: FieldImpl> {
-    width: u32,
-
-    alpha: u32,
-
-    internal_rounds: u32,
-
-    external_rounds: u32,
-
-    round_constants: &'a DeviceSlice<F>,
-
-    inernal_matrix_diag: &'a DeviceSlice<F>,
-
-    pub mds_type: MdsType,
-
-    pub diffusion: DiffusionStrategy,
-}
-
-impl<F: FieldImpl> std::fmt::Debug for Poseidon2Constants<'_, F> {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        write!(
-            f,
-            "{}, {}, {}, {}",
-            self.width, self.alpha, self.internal_rounds, self.external_rounds
-        )
+impl<F> Poseidon2<F>
+where
+    F: FieldImpl,
+    <F as FieldImpl>::Config: Poseidon2Impl<F>,
+{
+    pub fn load(
+        width: usize,
+        rate: usize,
+        mds_type: MdsType,
+        diffusion: DiffusionStrategy,
+        ctx: &DeviceContext,
+    ) -> IcicleResult<Self> {
+        <<F as FieldImpl>::Config as Poseidon2Impl<F>>::load(width as u32, rate as u32, mds_type, diffusion, ctx)
+            .and_then(|handle| {
+                Ok(Self {
+                    width,
+                    handle,
+                    phantom: PhantomData,
+                })
+            })
     }
-}
 
-/// Struct that encodes Poseidon parameters to be passed into the [poseidon_hash_many](poseidon_hash_many) function.
-#[repr(C)]
-#[derive(Debug, Clone)]
-pub struct Poseidon2Config<'a> {
-    /// Details related to the device such as its id and stream id. See [DeviceContext](@ref device_context::DeviceContext).
-    pub ctx: DeviceContext<'a>,
-
-    are_states_on_device: bool,
-
-    are_outputs_on_device: bool,
-
-    pub mode: PoseidonMode,
-
-    pub output_index: u32,
-
-    /// Whether to run Poseidon asynchronously. If set to `true`, Poseidon will be non-blocking
-    /// and you'd need to synchronize it explicitly by running `cudaStreamSynchronize` or `cudaDeviceSynchronize`.
-    /// If set to `false`, Poseidon will block the current CPU thread.
-    pub is_async: bool,
-}
-
-impl<'a> Default for Poseidon2Config<'a> {
-    fn default() -> Self {
-        Self::default_for_device(DEFAULT_DEVICE_ID)
-    }
-}
-
-impl<'a> Poseidon2Config<'a> {
-    pub fn default_for_device(device_id: usize) -> Self {
-        Self {
-            ctx: DeviceContext::default_for_device(device_id),
-            are_states_on_device: false,
-            are_outputs_on_device: false,
-            mode: PoseidonMode::Compression,
-            output_index: 1,
-            is_async: false,
-        }
-    }
-}
-
-pub trait Poseidon2<F: FieldImpl> {
-    fn create_constants<'a>(
-        width: u32,
+    pub fn new(
+        width: usize,
+        rate: usize,
         alpha: u32,
         internal_rounds: u32,
         external_rounds: u32,
@@ -110,191 +69,122 @@ pub trait Poseidon2<F: FieldImpl> {
         mds_type: MdsType,
         diffusion: DiffusionStrategy,
         ctx: &DeviceContext,
-    ) -> IcicleResult<Poseidon2Constants<'a, F>>;
-    fn load_constants<'a>(
+    ) -> IcicleResult<Self> {
+        <<F as FieldImpl>::Config as Poseidon2Impl<F>>::create(
+            width as u32,
+            rate as u32,
+            alpha,
+            internal_rounds,
+            external_rounds,
+            round_constants,
+            internal_matrix_diag,
+            mds_type,
+            diffusion,
+            ctx,
+        )
+        .and_then(|handle| {
+            Ok(Self {
+                width,
+                handle,
+                phantom: PhantomData,
+            })
+        })
+    }
+}
+
+impl<F> SpongeHash<F, F> for Poseidon2<F>
+where
+    F: FieldImpl,
+    <F as FieldImpl>::Config: Poseidon2Impl<F>,
+{
+    fn get_handle(&self) -> *const c_void {
+        self.handle
+    }
+
+    fn hash_many(
+        &self,
+        inputs: &(impl HostOrDeviceSlice<F> + ?Sized),
+        output: &mut (impl HostOrDeviceSlice<F> + ?Sized),
+        number_of_states: usize,
+        input_block_len: usize,
+        output_len: usize,
+        cfg: &SpongeConfig,
+    ) -> IcicleResult<()> {
+        sponge_check_input(
+            inputs,
+            number_of_states,
+            input_block_len,
+            cfg.input_rate as usize,
+            &cfg.ctx,
+        );
+        sponge_check_outputs(output, number_of_states, output_len, self.width, false, &cfg.ctx);
+
+        let mut local_cfg = cfg.clone();
+        local_cfg.are_inputs_on_device = inputs.is_on_device();
+        local_cfg.are_outputs_on_device = output.is_on_device();
+
+        <<F as FieldImpl>::Config as Poseidon2Impl<F>>::hash_many(
+            inputs,
+            output,
+            number_of_states as u32,
+            input_block_len as u32,
+            output_len as u32,
+            self.handle,
+            &local_cfg,
+        )
+    }
+
+    fn default_config<'a>(&self) -> SpongeConfig<'a> {
+        let mut cfg = SpongeConfig::default();
+        cfg.input_rate = self.width as u32;
+        cfg.output_rate = self.width as u32;
+        cfg
+    }
+}
+
+impl<F> Drop for Poseidon2<F>
+where
+    F: FieldImpl,
+    <F as FieldImpl>::Config: Poseidon2Impl<F>,
+{
+    fn drop(&mut self) {
+        <<F as FieldImpl>::Config as Poseidon2Impl<F>>::delete(self.handle).unwrap();
+    }
+}
+
+pub trait Poseidon2Impl<F: FieldImpl> {
+    fn create(
         width: u32,
+        rate: u32,
+        alpha: u32,
+        internal_rounds: u32,
+        external_rounds: u32,
+        round_constants: &[F],
+        internal_matrix_diag: &[F],
         mds_type: MdsType,
         diffusion: DiffusionStrategy,
         ctx: &DeviceContext,
-    ) -> IcicleResult<Poseidon2Constants<'a, F>>;
-    fn poseidon_unchecked(
-        states: &(impl HostOrDeviceSlice<F> + ?Sized),
+    ) -> IcicleResult<Poseidon2Handle>;
+
+    fn load(
+        width: u32,
+        rate: u32,
+        mds_type: MdsType,
+        diffusion: DiffusionStrategy,
+        ctx: &DeviceContext,
+    ) -> IcicleResult<Poseidon2Handle>;
+
+    fn hash_many(
+        inputs: &(impl HostOrDeviceSlice<F> + ?Sized),
         output: &mut (impl HostOrDeviceSlice<F> + ?Sized),
         number_of_states: u32,
-        width: u32,
-        constants: &Poseidon2Constants<F>,
-        config: &Poseidon2Config,
+        input_block_len: u32,
+        output_len: u32,
+        poseidon: Poseidon2Handle,
+        cfg: &SpongeConfig,
     ) -> IcicleResult<()>;
-    fn poseidon_unchecked_inplace(
-        states: &mut (impl HostOrDeviceSlice<F> + ?Sized),
-        number_of_states: u32,
-        width: u32,
-        constants: &Poseidon2Constants<F>,
-        config: &Poseidon2Config,
-    ) -> IcicleResult<()>;
-    fn release_constants(constants: &Poseidon2Constants<F>, ctx: &DeviceContext) -> IcicleResult<()>;
-}
 
-/// Loads pre-calculated poseidon constants on the GPU.
-pub fn load_poseidon2_constants<'a, F>(
-    width: u32,
-    mds_type: MdsType,
-    diffusion: DiffusionStrategy,
-    ctx: &DeviceContext,
-) -> IcicleResult<Poseidon2Constants<'a, F>>
-where
-    F: FieldImpl,
-    <F as FieldImpl>::Config: Poseidon2<F>,
-{
-    <<F as FieldImpl>::Config as Poseidon2<F>>::load_constants(width, mds_type, diffusion, ctx)
-}
-
-/// Creates new instance of poseidon constants on the GPU.
-pub fn create_poseidon2_constants<'a, F>(
-    width: u32,
-    alpha: u32,
-    ctx: &DeviceContext,
-    internal_rounds: u32,
-    external_rounds: u32,
-    round_constants: &mut [F],
-    internal_matrix_diag: &mut [F],
-    mds_type: MdsType,
-    diffusion: DiffusionStrategy,
-) -> IcicleResult<Poseidon2Constants<'a, F>>
-where
-    F: FieldImpl,
-    <F as FieldImpl>::Config: Poseidon2<F>,
-{
-    <<F as FieldImpl>::Config as Poseidon2<F>>::create_constants(
-        width,
-        alpha,
-        internal_rounds,
-        external_rounds,
-        round_constants,
-        internal_matrix_diag,
-        mds_type,
-        diffusion,
-        ctx,
-    )
-}
-
-fn poseidon_checks<F>(
-    states: &(impl HostOrDeviceSlice<F> + ?Sized),
-    output: &(impl HostOrDeviceSlice<F> + ?Sized),
-    number_of_states: u32,
-    width: u32,
-    config: &Poseidon2Config,
-) where
-    F: FieldImpl,
-    <F as FieldImpl>::Config: Poseidon2<F>,
-{
-    if states.len() < (number_of_states * width) as usize {
-        panic!(
-            "input len is {}; but needs to be at least {}",
-            states.len(),
-            number_of_states * width
-        );
-    }
-    if output.len() < number_of_states as usize {
-        panic!(
-            "output len is {}; but needs to be at least {}",
-            output.len(),
-            number_of_states
-        );
-    }
-
-    let ctx_device_id = config
-        .ctx
-        .device_id;
-    if let Some(device_id) = states.device_id() {
-        assert_eq!(
-            device_id, ctx_device_id,
-            "Device ids in input and context are different"
-        );
-    }
-
-    if let Some(device_id) = output.device_id() {
-        assert_eq!(
-            device_id, ctx_device_id,
-            "Device ids in output and context are different"
-        );
-    }
-    check_device(ctx_device_id);
-}
-
-/// Computes the poseidon hashes for multiple preimages.
-///
-/// # Arguments
-///
-/// * `input` - a pointer to the input data. May point to a vector of preimages or a vector of states filled with preimages.
-///
-/// * `output` - a pointer to the output data. Must be at least of size [number_of_states](number_of_states)
-///
-/// * `number_of_states` - number of input blocks of size `arity`
-///
-/// * `arity` - the arity of the hash function (the size of 1 preimage)
-///
-/// * `constants` - Poseidon constants.
-///
-/// * `config` - config used to specify extra arguments of the Poseidon.
-pub fn poseidon2_hash_many<F>(
-    states: &(impl HostOrDeviceSlice<F> + ?Sized),
-    output: &mut (impl HostOrDeviceSlice<F> + ?Sized),
-    number_of_states: u32,
-    width: u32,
-    constants: &Poseidon2Constants<F>,
-    config: &Poseidon2Config,
-) -> IcicleResult<()>
-where
-    F: FieldImpl,
-    <F as FieldImpl>::Config: Poseidon2<F>,
-{
-    poseidon_checks(states, output, number_of_states, width, config);
-    let mut local_cfg = config.clone();
-    local_cfg.are_states_on_device = states.is_on_device();
-    local_cfg.are_outputs_on_device = output.is_on_device();
-
-    <<F as FieldImpl>::Config as Poseidon2<F>>::poseidon_unchecked(
-        states,
-        output,
-        number_of_states,
-        width,
-        constants,
-        &local_cfg,
-    )
-}
-
-pub fn poseidon2_hash_many_inplace<F>(
-    states: &mut (impl HostOrDeviceSlice<F> + ?Sized),
-    number_of_states: u32,
-    width: u32,
-    constants: &Poseidon2Constants<F>,
-    config: &Poseidon2Config,
-) -> IcicleResult<()>
-where
-    F: FieldImpl,
-    <F as FieldImpl>::Config: Poseidon2<F>,
-{
-    poseidon_checks(states, states, number_of_states, width, config);
-    let mut local_cfg = config.clone();
-    local_cfg.are_states_on_device = states.is_on_device();
-    local_cfg.are_outputs_on_device = states.is_on_device();
-
-    <<F as FieldImpl>::Config as Poseidon2<F>>::poseidon_unchecked_inplace(
-        states,
-        number_of_states,
-        width,
-        constants,
-        &local_cfg,
-    )
-}
-
-pub fn release_poseidon2_constants<'a, F>(constants: &Poseidon2Constants<F>, ctx: &DeviceContext) -> IcicleResult<()>
-where
-    F: FieldImpl,
-    <F as FieldImpl>::Config: Poseidon2<F>,
-{
-    <<F as FieldImpl>::Config as Poseidon2<F>>::release_constants(constants, ctx)
+    fn delete(poseidon: Poseidon2Handle) -> IcicleResult<()>;
 }
 
 #[macro_export]
@@ -307,140 +197,125 @@ macro_rules! impl_poseidon2 {
     ) => {
         mod $field_prefix_ident {
             use crate::poseidon2::{
-                $field, $field_config, CudaError, DeviceContext, DiffusionStrategy, MdsType, Poseidon2Config,
-                Poseidon2Constants,
+                $field, $field_config, CudaError, DeviceContext, DiffusionStrategy, MdsType, Poseidon2Handle,
+                SpongeConfig,
             };
+            use icicle_core::error::IcicleError;
             extern "C" {
-                #[link_name = concat!($field_prefix, "_create_poseidon2_constants_cuda")]
-                pub(crate) fn _create_constants(
+                #[link_name = concat!($field_prefix, "_poseidon2_create_cuda")]
+                pub(crate) fn create(
+                    poseidon: *mut Poseidon2Handle,
                     width: u32,
+                    rate: u32,
                     alpha: u32,
                     internal_rounds: u32,
                     external_rounds: u32,
-                    constants: *mut $field,
-                    internal_matrix_diag: *mut $field,
+                    constants: *const $field,
+                    internal_matrix_diag: *const $field,
                     mds_type: MdsType,
                     diffusion: DiffusionStrategy,
                     ctx: &DeviceContext,
-                    poseidon_constants: *mut Poseidon2Constants<$field>,
                 ) -> CudaError;
 
-                #[link_name = concat!($field_prefix, "_init_poseidon2_constants_cuda")]
-                pub(crate) fn _load_constants(
+                #[link_name = concat!($field_prefix, "_poseidon2_load_cuda")]
+                pub(crate) fn load(
+                    poseidon: *mut Poseidon2Handle,
                     width: u32,
+                    rate: u32,
                     mds_type: MdsType,
                     diffusion: DiffusionStrategy,
                     ctx: &DeviceContext,
-                    constants: *mut Poseidon2Constants<$field>,
                 ) -> CudaError;
 
-                #[link_name = concat!($field_prefix, "_release_poseidon2_constants_cuda")]
-                pub(crate) fn _release_constants(
-                    constants: &Poseidon2Constants<$field>,
-                    ctx: &DeviceContext,
-                ) -> CudaError;
+                #[link_name = concat!($field_prefix, "_poseidon2_delete_cuda")]
+                pub(crate) fn delete(poseidon: Poseidon2Handle) -> CudaError;
 
-                #[link_name = concat!($field_prefix, "_poseidon2_hash_cuda")]
+                #[link_name = concat!($field_prefix, "_poseidon2_hash_many_cuda")]
                 pub(crate) fn hash_many(
-                    states: *const $field,
+                    poseidon: Poseidon2Handle,
+                    inputs: *const $field,
                     output: *mut $field,
                     number_of_states: u32,
-                    width: u32,
-                    constants: &Poseidon2Constants<$field>,
-                    config: &Poseidon2Config,
+                    input_block_len: u32,
+                    output_len: u32,
+                    cfg: &SpongeConfig,
                 ) -> CudaError;
             }
         }
 
-        impl Poseidon2<$field> for $field_config {
-            fn create_constants<'a>(
+        impl Poseidon2Impl<$field> for $field_config {
+            fn create(
                 width: u32,
+                rate: u32,
                 alpha: u32,
                 internal_rounds: u32,
                 external_rounds: u32,
-                round_constants: &mut [$field],
-                internal_matrix_diag: &mut [$field],
+                round_constants: &[$field],
+                internal_matrix_diag: &[$field],
                 mds_type: MdsType,
                 diffusion: DiffusionStrategy,
                 ctx: &DeviceContext,
-            ) -> IcicleResult<Poseidon2Constants<'a, $field>> {
+            ) -> IcicleResult<Poseidon2Handle> {
                 unsafe {
-                    let mut poseidon_constants = MaybeUninit::<Poseidon2Constants<'a, $field>>::uninit();
-                    let err = $field_prefix_ident::_create_constants(
+                    let mut poseidon = MaybeUninit::<Poseidon2Handle>::uninit();
+                    $field_prefix_ident::create(
+                        poseidon.as_mut_ptr(),
                         width,
+                        rate,
                         alpha,
                         internal_rounds,
                         external_rounds,
-                        round_constants as *mut _ as *mut $field,
-                        internal_matrix_diag as *mut _ as *mut $field,
+                        round_constants as *const _ as *const $field,
+                        internal_matrix_diag as *const _ as *const $field,
                         mds_type,
                         diffusion,
                         ctx,
-                        poseidon_constants.as_mut_ptr(),
                     )
-                    .wrap();
-                    err.and(Ok(poseidon_constants.assume_init()))
+                    .wrap()
+                    .and(Ok(poseidon.assume_init()))
                 }
             }
 
-            fn load_constants<'a>(
+            fn load(
                 width: u32,
+                rate: u32,
                 mds_type: MdsType,
                 diffusion: DiffusionStrategy,
                 ctx: &DeviceContext,
-            ) -> IcicleResult<Poseidon2Constants<'a, $field>> {
+            ) -> IcicleResult<Poseidon2Handle> {
                 unsafe {
-                    let mut constants = MaybeUninit::<Poseidon2Constants<'a, $field>>::uninit();
-                    let err =
-                        $field_prefix_ident::_load_constants(width, mds_type, diffusion, ctx, constants.as_mut_ptr())
-                            .wrap();
-                    err.and(Ok(constants.assume_init()))
+                    let mut poseidon = MaybeUninit::<Poseidon2Handle>::uninit();
+                    $field_prefix_ident::load(poseidon.as_mut_ptr(), width, rate, mds_type, diffusion, ctx)
+                        .wrap()
+                        .and(Ok(poseidon.assume_init()))
                 }
             }
 
-            fn poseidon_unchecked(
-                states: &(impl HostOrDeviceSlice<$field> + ?Sized),
+            fn hash_many(
+                inputs: &(impl HostOrDeviceSlice<$field> + ?Sized),
                 output: &mut (impl HostOrDeviceSlice<$field> + ?Sized),
                 number_of_states: u32,
-                width: u32,
-                constants: &Poseidon2Constants<$field>,
-                config: &Poseidon2Config,
+                input_block_len: u32,
+                output_len: u32,
+                poseidon: Poseidon2Handle,
+                cfg: &SpongeConfig,
             ) -> IcicleResult<()> {
                 unsafe {
                     $field_prefix_ident::hash_many(
-                        states.as_ptr(),
+                        poseidon,
+                        inputs.as_ptr(),
                         output.as_mut_ptr(),
                         number_of_states,
-                        width,
-                        constants,
-                        config,
+                        input_block_len,
+                        output_len,
+                        cfg,
                     )
                     .wrap()
                 }
             }
 
-            fn poseidon_unchecked_inplace(
-                states: &mut (impl HostOrDeviceSlice<$field> + ?Sized),
-                number_of_states: u32,
-                width: u32,
-                constants: &Poseidon2Constants<$field>,
-                config: &Poseidon2Config,
-            ) -> IcicleResult<()> {
-                unsafe {
-                    $field_prefix_ident::hash_many(
-                        states.as_ptr(),
-                        states.as_mut_ptr(),
-                        number_of_states,
-                        width,
-                        constants,
-                        config,
-                    )
-                    .wrap()
-                }
-            }
-
-            fn release_constants<'a>(constants: &Poseidon2Constants<$field>, ctx: &DeviceContext) -> IcicleResult<()> {
-                unsafe { $field_prefix_ident::_release_constants(constants, ctx).wrap() }
+            fn delete(poseidon: Poseidon2Handle) -> IcicleResult<()> {
+                unsafe { $field_prefix_ident::delete(poseidon).wrap() }
             }
         }
     };
@@ -466,42 +341,41 @@ pub mod bench {
     };
 
     use crate::{
+        hash::SpongeHash,
         ntt::FieldImpl,
-        poseidon2::{load_poseidon2_constants, DiffusionStrategy, MdsType},
+        poseidon2::{DiffusionStrategy, MdsType, Poseidon2, Poseidon2Impl},
         traits::GenerateRandom,
         vec_ops::VecOps,
     };
 
-    use super::{poseidon2_hash_many, Poseidon2, Poseidon2Config, Poseidon2Constants};
-
     #[allow(unused)]
-    fn poseidon2_for_bench<'a, F: FieldImpl>(
+    fn poseidon2_for_bench<F: FieldImpl>(
+        poseidon: &Poseidon2<F>,
         states: &(impl HostOrDeviceSlice<F> + ?Sized),
         poseidon2_result: &mut (impl HostOrDeviceSlice<F> + ?Sized),
         number_of_states: usize,
-        width: usize,
-        constants: &Poseidon2Constants<'a, F>,
-        config: &Poseidon2Config,
+        ctx: &DeviceContext,
         _seed: u32,
     ) where
-        <F as FieldImpl>::Config: Poseidon2<F> + GenerateRandom<F>,
-        <F as FieldImpl>::Config: VecOps<F>,
+        <F as FieldImpl>::Config: Poseidon2Impl<F> + GenerateRandom<F>,
     {
-        poseidon2_hash_many(
-            states,
-            poseidon2_result,
-            number_of_states as u32,
-            width as u32,
-            constants,
-            config,
-        )
-        .unwrap();
+        let cfg = poseidon.default_config();
+        poseidon
+            .hash_many(
+                states,
+                poseidon2_result,
+                number_of_states,
+                poseidon.width,
+                poseidon.width,
+                &cfg,
+            )
+            .unwrap();
     }
 
     #[allow(unused)]
     pub fn benchmark_poseidon2<F: FieldImpl>(c: &mut Criterion)
     where
-        <F as FieldImpl>::Config: Poseidon2<F> + GenerateRandom<F>,
+        <F as FieldImpl>::Config: Poseidon2Impl<F> + GenerateRandom<F>,
         <F as FieldImpl>::Config: VecOps<F>,
     {
         use criterion::SamplingMode;
@@ -519,7 +393,7 @@ pub mod bench {
             .parse::<u32>()
             .unwrap_or(MAX_LOG2);
 
-        for test_size_log2 in 13u32..max_log2 + 1 {
+        for test_size_log2 in 18u32..max_log2 + 1 {
             for t in [2, 3, 4, 8, 16, 20, 24] {
                 let number_of_states = 1 << test_size_log2;
                 let full_size = t * number_of_states;
@@ -531,31 +405,27 @@ pub mod bench {
                 let permutation_result_slice = HostSlice::from_mut_slice(&mut permutation_result);
 
                 let ctx = DeviceContext::default();
-                let config = Poseidon2Config::default();
-                for mds in [MdsType::Default, MdsType::Plonky] {
-                    for diffusion in [DiffusionStrategy::Default, DiffusionStrategy::Montgomery] {
-                        let constants =
-                            load_poseidon2_constants(t as u32, mds.clone(), diffusion.clone(), &ctx).unwrap();
-                        let bench_descr = format!(
-                            "Mds::{:?}; Diffusion::{:?}; Number of states: {}; Width: {}",
-                            mds, diffusion, number_of_states, t
-                        );
-                        group.bench_function(&bench_descr, |b| {
-                            b.iter(|| {
-                                poseidon2_for_bench::<F>(
-                                    input,
-                                    permutation_result_slice,
-                                    number_of_states,
-                                    t,
-                                    &constants,
-                                    &config,
-                                    black_box(1),
-                                )
-                            })
-                        });
-
-                        // }
-                    }
+                for (mds, diffusion) in [
+                    (MdsType::Default, DiffusionStrategy::Default),
+                    (MdsType::Plonky, DiffusionStrategy::Montgomery),
+                ] {
+                    let poseidon = Poseidon2::<F>::load(t, t, mds, diffusion, &ctx).unwrap();
+                    let bench_descr = format!(
+                        "TestSize: 2**{}, Mds::{:?}, Diffusion::{:?}, Width: {}",
+                        test_size_log2, mds, diffusion, t
+                    );
+                    group.bench_function(&bench_descr, |b| {
+                        b.iter(|| {
+                            poseidon2_for_bench::<F>(
+                                &poseidon,
+                                input,
+                                permutation_result_slice,
+                                number_of_states,
+                                &ctx,
+                                black_box(1),
+                            )
+                        })
+                    });
                 }
             }
         }

wrappers/rust/icicle-core/src/poseidon2/tests.rs

@@ -1,27 +1,21 @@
-use crate::poseidon2::{MdsType, PoseidonMode};
+use crate::hash::SpongeHash;
 use crate::traits::FieldImpl;
 use icicle_cuda_runtime::device_context::DeviceContext;
 use icicle_cuda_runtime::memory::{HostOrDeviceSlice, HostSlice};
 
-use super::{
-    load_poseidon2_constants, poseidon2_hash_many, DiffusionStrategy, Poseidon2, Poseidon2Config, Poseidon2Constants,
-};
+use super::{DiffusionStrategy, MdsType, Poseidon2, Poseidon2Impl};
 
-pub fn init_poseidon<'a, F: FieldImpl>(
-    width: u32,
-    mds_type: MdsType,
-    diffusion: DiffusionStrategy,
-) -> Poseidon2Constants<'a, F>
+pub fn init_poseidon<F: FieldImpl>(width: usize, mds_type: MdsType, diffusion: DiffusionStrategy) -> Poseidon2<F>
 where
-    <F as FieldImpl>::Config: Poseidon2<F>,
+    <F as FieldImpl>::Config: Poseidon2Impl<F>,
 {
     let ctx = DeviceContext::default();
-    load_poseidon2_constants::<F>(width, mds_type, diffusion, &ctx).unwrap()
+    Poseidon2::load(width, width, mds_type, diffusion, &ctx).unwrap()
 }
 
-fn _check_poseidon_hash_many<F: FieldImpl>(width: u32, constants: Poseidon2Constants<F>) -> (F, F)
+fn _check_poseidon_hash_many<F: FieldImpl>(width: usize, poseidon: &Poseidon2<F>) -> (F, F)
 where
-    <F as FieldImpl>::Config: Poseidon2<F>,
+    <F as FieldImpl>::Config: Poseidon2Impl<F>,
 {
     let test_size = 1 << 10;
     let mut inputs = vec![F::one(); test_size * width as usize];
@@ -30,16 +24,10 @@ where
     let input_slice = HostSlice::from_mut_slice(&mut inputs);
     let output_slice = HostSlice::from_mut_slice(&mut outputs);
 
-    let config = Poseidon2Config::default();
-    poseidon2_hash_many::<F>(
-        input_slice,
-        output_slice,
-        test_size as u32,
-        width as u32,
-        &constants,
-        &config,
-    )
-    .unwrap();
+    let cfg = poseidon.default_config();
+    poseidon
+        .hash_many(input_slice, output_slice, test_size, width, 1, &cfg)
+        .unwrap();
 
     let a1 = output_slice[0];
     let a2 = output_slice[output_slice.len() - 2];
@@ -49,21 +37,22 @@ where
     (a1, a2)
 }
 
-pub fn check_poseidon_hash_many<'a, F: FieldImpl + 'a>()
+pub fn check_poseidon_hash_many<F: FieldImpl>()
 where
-    <F as FieldImpl>::Config: Poseidon2<F>,
+    <F as FieldImpl>::Config: Poseidon2Impl<F>,
 {
     let widths = [2, 3, 4, 8, 12, 16, 20, 24];
+    let ctx = DeviceContext::default();
     for width in widths {
-        let constants = init_poseidon::<'a, F>(width as u32, MdsType::Default, DiffusionStrategy::Default);
+        let poseidon = Poseidon2::<F>::load(width, width, MdsType::Default, DiffusionStrategy::Default, &ctx).unwrap();
 
-        _check_poseidon_hash_many(width, constants);
+        _check_poseidon_hash_many(width, &poseidon);
     }
 }
 
-pub fn check_poseidon_kats<'a, F: FieldImpl>(width: usize, kats: &[F], constants: &Poseidon2Constants<'a, F>)
+pub fn check_poseidon_kats<F: FieldImpl>(width: usize, kats: &[F], poseidon: &Poseidon2<F>)
 where
-    <F as FieldImpl>::Config: Poseidon2<F>,
+    <F as FieldImpl>::Config: Poseidon2Impl<F>,
 {
     assert_eq!(width, kats.len());
 
@@ -83,17 +72,11 @@ where
     let input_slice = HostSlice::from_mut_slice(&mut inputs);
     let output_slice = HostSlice::from_mut_slice(&mut outputs);
 
-    let mut config = Poseidon2Config::default();
-    config.mode = PoseidonMode::Permutation;
-    poseidon2_hash_many::<F>(
-        input_slice,
-        output_slice,
-        batch_size as u32,
-        width as u32,
-        &constants,
-        &config,
-    )
-    .unwrap();
+    let cfg = poseidon.default_config();
+
+    poseidon
+        .hash_many(input_slice, output_slice, batch_size, width, width, &cfg)
+        .unwrap();
 
     for (i, val) in output_slice
         .iter()

wrappers/rust/icicle-core/src/tree/mmcs.rs

@@ -0,0 +1,79 @@
+use icicle_cuda_runtime::memory::HostSlice;
+
+use crate::{error::IcicleResult, ntt::FieldImpl};
+use crate::{hash::SpongeHash, Matrix};
+
+use super::TreeBuilderConfig;
+
+pub trait FieldMmcs<F, Compression, Hasher>
+where
+    F: FieldImpl,
+    Compression: SpongeHash<F, F>,
+    Hasher: SpongeHash<F, F>,
+{
+    fn mmcs_commit(
+        leaves: Vec<Matrix>,
+        digests: &mut HostSlice<F>,
+        hasher: &Hasher,
+        compression: &Compression,
+        config: &TreeBuilderConfig,
+    ) -> IcicleResult<()>;
+}
+
+#[macro_export]
+macro_rules! impl_mmcs {
+    (
+      $field_prefix:literal,
+      $field_prefix_ident:ident,
+      $field:ident,
+      $field_config:ident,
+      $mmcs:ident
+    ) => {
+        mod $field_prefix_ident {
+            use super::*;
+            use icicle_cuda_runtime::error::CudaError;
+
+            extern "C" {
+                #[link_name = concat!($field_prefix, "_mmcs_commit_cuda")]
+                pub(crate) fn mmcs_commit_cuda(
+                    leaves: *const Matrix,
+                    number_of_inputs: u32,
+                    digests: *mut $field,
+                    hasher: *const c_void,
+                    compression: *const c_void,
+                    config: &TreeBuilderConfig,
+                ) -> CudaError;
+            }
+        }
+
+        struct $mmcs;
+
+        impl<Compression, Hasher> FieldMmcs<$field, Compression, Hasher> for $mmcs
+        where
+            Compression: SpongeHash<$field, $field>,
+            Hasher: SpongeHash<$field, $field>,
+        {
+            fn mmcs_commit(
+                leaves: Vec<Matrix>,
+                digests: &mut HostSlice<$field>,
+                hasher: &Hasher,
+                compression: &Compression,
+                config: &TreeBuilderConfig,
+            ) -> IcicleResult<()> {
+                unsafe {
+                    $field_prefix_ident::mmcs_commit_cuda(
+                        leaves
+                            .as_slice()
+                            .as_ptr(),
+                        leaves.len() as u32,
+                        digests.as_mut_ptr(),
+                        compression.get_handle(),
+                        hasher.get_handle(),
+                        config,
+                    )
+                    .wrap()
+                }
+            }
+        }
+    };
+}

wrappers/rust/icicle-core/src/tree/mod.rs

@@ -1,11 +1,12 @@
-use icicle_cuda_runtime::device::check_device;
 use icicle_cuda_runtime::{
     device_context::{DeviceContext, DEFAULT_DEVICE_ID},
     memory::HostOrDeviceSlice,
 };
 
-use crate::{error::IcicleResult, poseidon::PoseidonConstants, traits::FieldImpl};
+use crate::hash::SpongeHash;
+use crate::{error::IcicleResult, ntt::FieldImpl};
 
+pub mod mmcs;
 #[doc(hidden)]
 pub mod tests;
 
@@ -16,11 +17,20 @@ pub struct TreeBuilderConfig<'a> {
     /// Details related to the device such as its id and stream id. See [DeviceContext](@ref device_context::DeviceContext).
     pub ctx: DeviceContext<'a>,
 
+    /// Airty of the tree
+    pub arity: u32,
+
     /// How many rows of the Merkle tree rows should be written to output. '0' means all of them
-    keep_rows: u32,
+    pub keep_rows: u32,
+
+    /// The size of the output for the bottom layer hash and compression.
+    /// Will also be equal to the size of the root of the tree. Default value 1
+    pub digest_elements: u32,
 
     are_inputs_on_device: bool,
 
+    are_outputs_on_device: bool,
+
     /// Whether to run build_merkle_tree asynchronously. If set to `true`, TreeBuilder will be non-blocking
     /// and you'd need to synchronize it explicitly by running `cudaStreamSynchronize` or `cudaDeviceSynchronize`.
     /// If set to `false`, build_merkle_tree will block the current CPU thread.
@@ -32,134 +42,99 @@ impl<'a> Default for TreeBuilderConfig<'a> {
         Self::default_for_device(DEFAULT_DEVICE_ID)
     }
 }
+
 impl<'a> TreeBuilderConfig<'a> {
     fn default_for_device(device_id: usize) -> Self {
         Self {
             ctx: DeviceContext::default_for_device(device_id),
+            arity: 2,
             keep_rows: 0,
+            digest_elements: 1,
             are_inputs_on_device: false,
+            are_outputs_on_device: false,
             is_async: false,
         }
     }
 }
 
-pub fn merkle_tree_digests_len(height: u32, arity: u32) -> usize {
+pub fn merkle_tree_digests_len(height: u32, arity: u32, digest_elements: u32) -> usize {
     let mut digests_len = 0usize;
-    let mut row_length = 1;
-    for _ in 1..height {
+    let mut row_length = digest_elements as usize;
+    for _ in 0..height + 1 {
         digests_len += row_length;
         row_length *= arity as usize;
     }
     digests_len
 }
 
-pub trait TreeBuilder<F: FieldImpl> {
-    fn build_poseidon_tree_unchecked(
-        leaves: &mut (impl HostOrDeviceSlice<F> + ?Sized),
-        digests: &mut [F],
-        height: u32,
-        arity: u32,
-        constants: &PoseidonConstants<F>,
+pub trait FieldTreeBuilder<F, Compression, Sponge>
+where
+    F: FieldImpl,
+    Compression: SpongeHash<F, F>,
+    Sponge: SpongeHash<F, F>,
+{
+    fn build_merkle_tree(
+        leaves: &(impl HostOrDeviceSlice<F> + ?Sized),
+        digests: &mut (impl HostOrDeviceSlice<F> + ?Sized),
+        height: usize,
+        input_block_len: usize,
+        compression: &Compression,
+        sponge: &Sponge,
         config: &TreeBuilderConfig,
     ) -> IcicleResult<()>;
 }
 
-/// Builds a Poseidon Merkle tree.
-///
-/// # Arguments
-///
-/// * `leaves` - a pointer to the leaves layer. Expected to have arity ^ (height - 1) elements
-///
-/// * `digests` - a pointer to the digests storage. Expected to have `sum(arity ^ (i)) for i in [0..height-1]`
-///
-/// * `height` - the height of the merkle tree
-///
-/// * `config` - config used to specify extra arguments of the Tree builder.
-pub fn build_poseidon_merkle_tree<F>(
-    leaves: &mut (impl HostOrDeviceSlice<F> + ?Sized),
-    digests: &mut [F],
-    height: u32,
-    arity: u32,
-    constants: &PoseidonConstants<F>,
-    config: &TreeBuilderConfig,
-) -> IcicleResult<()>
-where
-    F: FieldImpl,
-    <F as FieldImpl>::Config: TreeBuilder<F>,
-{
-    let leaves_len = 1 << (height - 1) as usize;
-    if leaves.len() != leaves_len {
-        panic!("Leaves len is {}; but needs to be exactly {}", leaves.len(), leaves_len,);
-    }
-
-    let digests_len = merkle_tree_digests_len(height, arity);
-    if digests.len() != digests_len as usize {
-        panic!(
-            "Digests len is {}; but needs to be exactly {}",
-            digests.len(),
-            digests_len
-        );
-    }
-
-    let ctx_device_id = config
-        .ctx
-        .device_id;
-    if let Some(device_id) = leaves.device_id() {
-        assert_eq!(
-            device_id, ctx_device_id,
-            "Device ids in leaves and context are different"
-        );
-    }
-    check_device(ctx_device_id);
-    let mut local_cfg = config.clone();
-    local_cfg.are_inputs_on_device = leaves.is_on_device();
-
-    <<F as FieldImpl>::Config as TreeBuilder<F>>::build_poseidon_tree_unchecked(
-        leaves, digests, height, arity, constants, &local_cfg,
-    )
-}
-
 #[macro_export]
-macro_rules! impl_tree_builder {
+macro_rules! impl_field_tree_builder {
     (
       $field_prefix:literal,
       $field_prefix_ident:ident,
       $field:ident,
-      $field_config:ident
+      $field_config:ident,
+      $tree_builder:ident
     ) => {
         mod $field_prefix_ident {
-            use crate::tree::{$field, $field_config, CudaError, DeviceContext, TreeBuilderConfig};
-            use icicle_core::poseidon::PoseidonConstants;
+            use super::*;
+            use icicle_cuda_runtime::error::CudaError;
 
             extern "C" {
-                #[link_name = concat!($field_prefix, "_build_poseidon_merkle_tree")]
-                pub(crate) fn _build_poseidon_merkle_tree(
-                    leaves: *mut $field,
+                #[link_name = concat!($field_prefix, "_build_merkle_tree")]
+                pub(crate) fn build_merkle_tree(
+                    leaves: *const $field,
                     digests: *mut $field,
                     height: u32,
-                    arity: u32,
-                    constants: &PoseidonConstants<$field>,
+                    input_block_len: u32,
+                    compression: *const c_void,
+                    sponge: *const c_void,
                     config: &TreeBuilderConfig,
                 ) -> CudaError;
             }
         }
 
-        impl TreeBuilder<$field> for $field_config {
-            fn build_poseidon_tree_unchecked(
-                leaves: &mut (impl HostOrDeviceSlice<$field> + ?Sized),
-                digests: &mut [$field],
-                height: u32,
-                arity: u32,
-                constants: &PoseidonConstants<$field>,
+        struct $tree_builder;
+
+        impl<Compression, Sponge> FieldTreeBuilder<$field, Compression, Sponge> for $tree_builder
+        where
+            Compression: SpongeHash<$field, $field>,
+            Sponge: SpongeHash<$field, $field>,
+        {
+            fn build_merkle_tree(
+                leaves: &(impl HostOrDeviceSlice<$field> + ?Sized),
+                digests: &mut (impl HostOrDeviceSlice<$field> + ?Sized),
+                height: usize,
+                input_block_len: usize,
+                compression: &Compression,
+                sponge: &Sponge,
                 config: &TreeBuilderConfig,
             ) -> IcicleResult<()> {
                 unsafe {
-                    $field_prefix_ident::_build_poseidon_merkle_tree(
-                        leaves.as_mut_ptr(),
-                        digests as *mut _ as *mut $field,
-                        height,
-                        arity,
-                        constants,
+                    $field_prefix_ident::build_merkle_tree(
+                        leaves.as_ptr(),
+                        digests.as_mut_ptr(),
+                        height as u32,
+                        input_block_len as u32,
+                        compression.get_handle(),
+                        sponge.get_handle(),
                         config,
                     )
                     .wrap()
@@ -168,15 +143,3 @@ macro_rules! impl_tree_builder {
         }
     };
 }
-
-#[macro_export]
-macro_rules! impl_tree_builder_tests {
-    (
-      $field:ident
-    ) => {
-        #[test]
-        fn test_build_poseidon_merkle_tree() {
-            check_build_merkle_tree::<$field>()
-        }
-    };
-}

wrappers/rust/icicle-core/src/tree/tests.rs

@@ -1,30 +1,42 @@
 use icicle_cuda_runtime::memory::HostSlice;
 
 use crate::{
-    poseidon::{tests::init_poseidon, Poseidon},
+    hash::SpongeHash,
     traits::FieldImpl,
-    tree::{build_poseidon_merkle_tree, merkle_tree_digests_len, TreeBuilderConfig},
+    tree::{merkle_tree_digests_len, TreeBuilderConfig},
 };
 
-use super::TreeBuilder;
+use super::FieldTreeBuilder;
 
-pub fn check_build_merkle_tree<F: FieldImpl>()
-where
-    <F as FieldImpl>::Config: TreeBuilder<F> + Poseidon<F>,
+pub fn check_build_field_merkle_tree<F, H, T>(
+    height: usize,
+    arity: usize,
+    sponge: &H,
+    compression: &H,
+    _expected_root: F,
+) where
+    F: FieldImpl,
+    H: SpongeHash<F, F>,
+    T: FieldTreeBuilder<F, H, H>,
 {
-    let height = 20;
-    let arity = 2;
-    let keep_rows = 1;
-    let mut leaves = vec![F::one(); 1 << (height - 1)];
-    let mut digests = vec![F::zero(); merkle_tree_digests_len(height, arity)];
-
-    let leaves_slice = HostSlice::from_mut_slice(&mut leaves);
-
-    let constants = init_poseidon(arity as u32);
-
     let mut config = TreeBuilderConfig::default();
-    config.keep_rows = keep_rows;
-    build_poseidon_merkle_tree::<F>(leaves_slice, &mut digests, height, arity, &constants, &config).unwrap();
+    config.arity = arity as u32;
+    let input_block_len = arity;
+    let leaves = vec![F::one(); (1 << height) * arity];
+    let mut digests = vec![F::zero(); merkle_tree_digests_len((height + 1) as u32, arity as u32, 1)];
 
-    println!("Root: {:?}", digests[0]);
+    let leaves_slice = HostSlice::from_slice(&leaves);
+    let digests_slice = HostSlice::from_mut_slice(&mut digests);
+
+    T::build_merkle_tree(
+        leaves_slice,
+        digests_slice,
+        height,
+        input_block_len,
+        compression,
+        sponge,
+        &config,
+    )
+    .unwrap();
+    println!("Root: {:?}", digests_slice[0]);
 }

wrappers/rust/icicle-curves/icicle-bls12-377/src/poseidon/mod.rs

@@ -3,8 +3,9 @@ use crate::curve::{BaseCfg, BaseField};
 use crate::curve::{ScalarCfg, ScalarField};
 
 use icicle_core::error::IcicleResult;
+use icicle_core::hash::SpongeConfig;
 use icicle_core::impl_poseidon;
-use icicle_core::poseidon::{Poseidon, PoseidonConfig, PoseidonConstants};
+use icicle_core::poseidon::{PoseidonHandle, PoseidonImpl};
 use icicle_core::traits::IcicleResultWrap;
 use icicle_cuda_runtime::device_context::DeviceContext;
 use icicle_cuda_runtime::error::CudaError;
@@ -20,9 +21,8 @@ impl_poseidon!("bw6_761", bw6_761, BaseField, BaseCfg);
 #[cfg(test)]
 pub(crate) mod tests {
     use crate::curve::ScalarField;
+    use icicle_core::impl_poseidon_tests;
     use icicle_core::poseidon::tests::*;
-    use icicle_core::{impl_poseidon_custom_config_test, impl_poseidon_tests};
 
     impl_poseidon_tests!(ScalarField);
-    impl_poseidon_custom_config_test!(ScalarField, 32, "bls12_377", 56);
 }

wrappers/rust/icicle-curves/icicle-bls12-377/src/tree/mod.rs

@@ -1,26 +1,29 @@
-#[cfg(feature = "bw6-761")]
-use crate::curve::{BaseCfg, BaseField};
-use crate::curve::{ScalarCfg, ScalarField};
-
 use icicle_core::error::IcicleResult;
-use icicle_core::impl_tree_builder;
-use icicle_core::poseidon::PoseidonConstants;
+use icicle_core::hash::SpongeHash;
+use icicle_core::impl_field_tree_builder;
 use icicle_core::traits::IcicleResultWrap;
-use icicle_core::tree::{TreeBuilder, TreeBuilderConfig};
-use icicle_cuda_runtime::device_context::DeviceContext;
-use icicle_cuda_runtime::error::CudaError;
+use icicle_core::tree::{FieldTreeBuilder, TreeBuilderConfig};
 use icicle_cuda_runtime::memory::HostOrDeviceSlice;
+use std::ffi::c_void;
 
-impl_tree_builder!("bls12_377", bls12_377, ScalarField, ScalarCfg);
+use crate::curve::ScalarField;
 
-#[cfg(feature = "bw6-761")]
-impl_tree_builder!("bw6_761", bw6_761, BaseField, BaseCfg);
+impl_field_tree_builder!("bls12_377", bls12_377_tb, ScalarField, ScalarCfg, Bls12_377TreeBuilder);
 
 #[cfg(test)]
 pub(crate) mod tests {
-    use crate::curve::ScalarField;
-    use icicle_core::impl_tree_builder_tests;
-    use icicle_core::tree::tests::*;
+    use icicle_core::{ntt::FieldImpl, poseidon::Poseidon, tree::tests::check_build_field_merkle_tree};
+    use icicle_cuda_runtime::device_context;
 
-    impl_tree_builder_tests!(ScalarField);
+    use crate::curve::ScalarField;
+
+    use super::Bls12_377TreeBuilder;
+
+    #[test]
+    fn poseidon_merkle_tree_test() {
+        let ctx = device_context::DeviceContext::default();
+        let sponge = Poseidon::load(2, &ctx).unwrap();
+
+        check_build_field_merkle_tree::<_, _, Bls12_377TreeBuilder>(25, 2, &sponge, &sponge, ScalarField::zero());
+    }
 }

wrappers/rust/icicle-curves/icicle-bls12-381/src/poseidon/mod.rs

@@ -1,8 +1,9 @@
 use crate::curve::{ScalarCfg, ScalarField};
 
 use icicle_core::error::IcicleResult;
+use icicle_core::hash::SpongeConfig;
 use icicle_core::impl_poseidon;
-use icicle_core::poseidon::{Poseidon, PoseidonConfig, PoseidonConstants};
+use icicle_core::poseidon::{PoseidonHandle, PoseidonImpl};
 use icicle_core::traits::IcicleResultWrap;
 use icicle_cuda_runtime::device_context::DeviceContext;
 use icicle_cuda_runtime::error::CudaError;
@@ -15,9 +16,8 @@ impl_poseidon!("bls12_381", bls12_381, ScalarField, ScalarCfg);
 #[cfg(test)]
 pub(crate) mod tests {
     use crate::curve::ScalarField;
+    use icicle_core::impl_poseidon_tests;
     use icicle_core::poseidon::tests::*;
-    use icicle_core::{impl_poseidon_custom_config_test, impl_poseidon_tests};
 
     impl_poseidon_tests!(ScalarField);
-    impl_poseidon_custom_config_test!(ScalarField, 32, "bls12_381", 55);
 }

wrappers/rust/icicle-curves/icicle-bls12-381/src/tree/mod.rs

@@ -1,21 +1,29 @@
-use crate::curve::{ScalarCfg, ScalarField};
-
 use icicle_core::error::IcicleResult;
-use icicle_core::impl_tree_builder;
-use icicle_core::poseidon::PoseidonConstants;
+use icicle_core::hash::SpongeHash;
+use icicle_core::impl_field_tree_builder;
 use icicle_core::traits::IcicleResultWrap;
-use icicle_core::tree::{TreeBuilder, TreeBuilderConfig};
-use icicle_cuda_runtime::device_context::DeviceContext;
-use icicle_cuda_runtime::error::CudaError;
+use icicle_core::tree::{FieldTreeBuilder, TreeBuilderConfig};
 use icicle_cuda_runtime::memory::HostOrDeviceSlice;
+use std::ffi::c_void;
 
-impl_tree_builder!("bls12_381", bls12_381, ScalarField, ScalarCfg);
+use crate::curve::ScalarField;
+
+impl_field_tree_builder!("bls12_381", bls12_381_tb, ScalarField, ScalarCfg, Bls12_381TreeBuilder);
 
 #[cfg(test)]
 pub(crate) mod tests {
-    use crate::curve::ScalarField;
-    use icicle_core::impl_tree_builder_tests;
-    use icicle_core::tree::tests::*;
+    use icicle_core::{ntt::FieldImpl, poseidon::Poseidon, tree::tests::check_build_field_merkle_tree};
+    use icicle_cuda_runtime::device_context;
 
-    impl_tree_builder_tests!(ScalarField);
+    use crate::curve::ScalarField;
+
+    use super::Bls12_381TreeBuilder;
+
+    #[test]
+    fn poseidon_merkle_tree_test() {
+        let ctx = device_context::DeviceContext::default();
+        let sponge = Poseidon::load(2, &ctx).unwrap();
+
+        check_build_field_merkle_tree::<_, _, Bls12_381TreeBuilder>(25, 2, &sponge, &sponge, ScalarField::zero());
+    }
 }

wrappers/rust/icicle-curves/icicle-bn254/src/poseidon/mod.rs

@@ -1,8 +1,9 @@
 use crate::curve::{ScalarCfg, ScalarField};
 
 use icicle_core::error::IcicleResult;
+use icicle_core::hash::SpongeConfig;
 use icicle_core::impl_poseidon;
-use icicle_core::poseidon::{Poseidon, PoseidonConfig, PoseidonConstants};
+use icicle_core::poseidon::{PoseidonHandle, PoseidonImpl};
 use icicle_core::traits::IcicleResultWrap;
 use icicle_cuda_runtime::device_context::DeviceContext;
 use icicle_cuda_runtime::error::CudaError;
@@ -15,9 +16,8 @@ impl_poseidon!("bn254", bn254, ScalarField, ScalarCfg);
 #[cfg(test)]
 pub(crate) mod tests {
     use crate::curve::ScalarField;
+    use icicle_core::impl_poseidon_tests;
     use icicle_core::poseidon::tests::*;
-    use icicle_core::{impl_poseidon_custom_config_test, impl_poseidon_tests};
 
     impl_poseidon_tests!(ScalarField);
-    impl_poseidon_custom_config_test!(ScalarField, 32, "bn254", 56);
 }

wrappers/rust/icicle-curves/icicle-bn254/src/poseidon2/mod.rs

@@ -1,8 +1,9 @@
 use crate::curve::{ScalarCfg, ScalarField};
 
 use icicle_core::error::IcicleResult;
+use icicle_core::hash::SpongeConfig;
 use icicle_core::impl_poseidon2;
-use icicle_core::poseidon2::{DiffusionStrategy, MdsType, Poseidon2, Poseidon2Config, Poseidon2Constants};
+use icicle_core::poseidon2::{DiffusionStrategy, MdsType, Poseidon2Handle, Poseidon2Impl};
 use icicle_core::traits::IcicleResultWrap;
 use icicle_cuda_runtime::device_context::DeviceContext;
 use icicle_cuda_runtime::error::CudaError;
@@ -29,7 +30,7 @@ pub(crate) mod tests {
             ScalarField::from_hex("0x1ed25194542b12eef8617361c3ba7c52e660b145994427cc86296242cf766ec8"),
         ];
 
-        let constants = init_poseidon::<ScalarField>(3, MdsType::Default, DiffusionStrategy::Default);
-        check_poseidon_kats(3, &kats, &constants);
+        let poseidon = init_poseidon::<ScalarField>(3, MdsType::Default, DiffusionStrategy::Default);
+        check_poseidon_kats(3, &kats, &poseidon);
     }
 }

wrappers/rust/icicle-curves/icicle-bn254/src/tree/mmcs.rs

@@ -0,0 +1,11 @@
+use icicle_core::error::IcicleResult;
+use icicle_core::hash::SpongeHash;
+use icicle_core::traits::IcicleResultWrap;
+use icicle_core::tree::{mmcs::FieldMmcs, TreeBuilderConfig};
+use icicle_core::{impl_mmcs, Matrix};
+use icicle_cuda_runtime::memory::{HostOrDeviceSlice, HostSlice};
+use std::ffi::c_void;
+
+use crate::curve::ScalarField;
+
+impl_mmcs!("bn254", bn254_mmcs, ScalarField, ScalarCfg, Bn254Mmcs);

wrappers/rust/icicle-curves/icicle-bn254/src/tree/mod.rs

@@ -1,21 +1,44 @@
-use crate::curve::{ScalarCfg, ScalarField};
+pub mod mmcs;
 
 use icicle_core::error::IcicleResult;
-use icicle_core::impl_tree_builder;
-use icicle_core::poseidon::PoseidonConstants;
+use icicle_core::hash::SpongeHash;
+use icicle_core::impl_field_tree_builder;
 use icicle_core::traits::IcicleResultWrap;
-use icicle_core::tree::{TreeBuilder, TreeBuilderConfig};
-use icicle_cuda_runtime::device_context::DeviceContext;
-use icicle_cuda_runtime::error::CudaError;
+use icicle_core::tree::{FieldTreeBuilder, TreeBuilderConfig};
 use icicle_cuda_runtime::memory::HostOrDeviceSlice;
+use std::ffi::c_void;
 
-impl_tree_builder!("bn254", bn254, ScalarField, ScalarCfg);
+use crate::curve::ScalarField;
+
+impl_field_tree_builder!("bn254", bn254_tb, ScalarField, ScalarCfg, Bn254TreeBuilder);
 
 #[cfg(test)]
 pub(crate) mod tests {
-    use crate::curve::ScalarField;
-    use icicle_core::impl_tree_builder_tests;
-    use icicle_core::tree::tests::*;
+    use icicle_core::{
+        ntt::FieldImpl,
+        poseidon::Poseidon,
+        poseidon2::{DiffusionStrategy, MdsType, Poseidon2},
+        tree::tests::check_build_field_merkle_tree,
+    };
+    use icicle_cuda_runtime::device_context;
 
-    impl_tree_builder_tests!(ScalarField);
+    use crate::curve::ScalarField;
+
+    use super::Bn254TreeBuilder;
+
+    #[test]
+    fn poseidon_merkle_tree_test() {
+        let ctx = device_context::DeviceContext::default();
+        let sponge = Poseidon::load(2, &ctx).unwrap();
+
+        check_build_field_merkle_tree::<_, _, Bn254TreeBuilder>(25, 2, &sponge, &sponge, ScalarField::zero());
+    }
+
+    #[test]
+    fn poseidon2_merkle_tree_test() {
+        let ctx = device_context::DeviceContext::default();
+        let sponge = Poseidon2::load(2, 2, MdsType::Default, DiffusionStrategy::Default, &ctx).unwrap();
+
+        check_build_field_merkle_tree::<_, _, Bn254TreeBuilder>(28, 2, &sponge, &sponge, ScalarField::zero());
+    }
 }

wrappers/rust/icicle-curves/icicle-bw6-761/src/lib.rs

@@ -2,7 +2,6 @@ pub mod curve;
 pub mod msm;
 pub mod ntt;
 pub mod poseidon;
-pub mod tree;
 pub mod vec_ops;
 
 impl icicle_core::SNARKCurve for curve::CurveCfg {}

wrappers/rust/icicle-curves/icicle-bw6-761/src/poseidon/mod.rs

@@ -1,9 +1,8 @@
 #[cfg(test)]
 pub(crate) mod tests {
     use crate::curve::ScalarField;
+    use icicle_core::impl_poseidon_tests;
     use icicle_core::poseidon::tests::*;
-    use icicle_core::{impl_poseidon_custom_config_test, impl_poseidon_tests};
 
     impl_poseidon_tests!(ScalarField);
-    impl_poseidon_custom_config_test!(ScalarField, 48, "bw6-761", 56);
 }

wrappers/rust/icicle-curves/icicle-bw6-761/src/tree/mod.rs

@@ -1,8 +0,0 @@
-#[cfg(test)]
-pub(crate) mod tests {
-    use crate::curve::ScalarField;
-    use icicle_core::impl_tree_builder_tests;
-    use icicle_core::tree::tests::*;
-
-    impl_tree_builder_tests!(ScalarField);
-}

wrappers/rust/icicle-curves/icicle-grumpkin/src/poseidon/mod.rs

@@ -1,8 +1,9 @@
 use crate::curve::{ScalarCfg, ScalarField};
 
 use icicle_core::error::IcicleResult;
+use icicle_core::hash::SpongeConfig;
 use icicle_core::impl_poseidon;
-use icicle_core::poseidon::{Poseidon, PoseidonConfig, PoseidonConstants};
+use icicle_core::poseidon::{PoseidonHandle, PoseidonImpl};
 use icicle_core::traits::IcicleResultWrap;
 use icicle_cuda_runtime::device_context::DeviceContext;
 use icicle_cuda_runtime::error::CudaError;
@@ -15,9 +16,8 @@ impl_poseidon!("grumpkin", grumpkin, ScalarField, ScalarCfg);
 #[cfg(test)]
 pub(crate) mod tests {
     use crate::curve::ScalarField;
+    use icicle_core::impl_poseidon_tests;
     use icicle_core::poseidon::tests::*;
-    use icicle_core::{impl_poseidon_custom_config_test, impl_poseidon_tests};
 
     impl_poseidon_tests!(ScalarField);
-    impl_poseidon_custom_config_test!(ScalarField, 32, "grumpkin", 56);
 }

wrappers/rust/icicle-curves/icicle-grumpkin/src/tree/mod.rs

@@ -1,21 +1,29 @@
-use crate::curve::{ScalarCfg, ScalarField};
-
 use icicle_core::error::IcicleResult;
-use icicle_core::impl_tree_builder;
-use icicle_core::poseidon::PoseidonConstants;
+use icicle_core::hash::SpongeHash;
+use icicle_core::impl_field_tree_builder;
 use icicle_core::traits::IcicleResultWrap;
-use icicle_core::tree::{TreeBuilder, TreeBuilderConfig};
-use icicle_cuda_runtime::device_context::DeviceContext;
-use icicle_cuda_runtime::error::CudaError;
+use icicle_core::tree::{FieldTreeBuilder, TreeBuilderConfig};
 use icicle_cuda_runtime::memory::HostOrDeviceSlice;
+use std::ffi::c_void;
 
-impl_tree_builder!("grumpkin", grumpkin, ScalarField, ScalarCfg);
+use crate::curve::ScalarField;
+
+impl_field_tree_builder!("grumpkin", grumpkin_tb, ScalarField, ScalarCfg, GrumpkinTreeBuilder);
 
 #[cfg(test)]
 pub(crate) mod tests {
-    use crate::curve::ScalarField;
-    use icicle_core::impl_tree_builder_tests;
-    use icicle_core::tree::tests::*;
+    use icicle_core::{ntt::FieldImpl, poseidon::Poseidon, tree::tests::check_build_field_merkle_tree};
+    use icicle_cuda_runtime::device_context;
 
-    impl_tree_builder_tests!(ScalarField);
+    use crate::curve::ScalarField;
+
+    use super::GrumpkinTreeBuilder;
+
+    #[test]
+    fn poseidon_merkle_tree_test() {
+        let ctx = device_context::DeviceContext::default();
+        let sponge = Poseidon::load(2, &ctx).unwrap();
+
+        check_build_field_merkle_tree::<_, _, GrumpkinTreeBuilder>(25, 2, &sponge, &sponge, ScalarField::zero());
+    }
 }

wrappers/rust/icicle-fields/icicle-babybear/Cargo.toml

@@ -18,6 +18,7 @@ cmake = "0.1.50"
 criterion = "0.3"
 risc0-core = "0.21.0"
 risc0-zkp = "0.21.0"
+
 p3-baby-bear = { git = "https://github.com/Plonky3/Plonky3", rev = "1e87146ebfaedc2150b635b10a096b733795fdce" }
 p3-symmetric = { git = "https://github.com/Plonky3/Plonky3", rev = "1e87146ebfaedc2150b635b10a096b733795fdce" }
 p3-mds = { git = "https://github.com/Plonky3/Plonky3", rev = "1e87146ebfaedc2150b635b10a096b733795fdce" }
@@ -25,6 +26,8 @@ p3-poseidon2 = { git = "https://github.com/Plonky3/Plonky3", rev = "1e87146ebfae
 p3-field = { git = "https://github.com/Plonky3/Plonky3", rev = "1e87146ebfaedc2150b635b10a096b733795fdce" }
 p3-dft = { git = "https://github.com/Plonky3/Plonky3", rev = "1e87146ebfaedc2150b635b10a096b733795fdce" }
 p3-matrix = { git = "https://github.com/Plonky3/Plonky3", rev = "1e87146ebfaedc2150b635b10a096b733795fdce" }
+p3-merkle-tree = { git = "https://github.com/Plonky3/Plonky3", rev = "1e87146ebfaedc2150b635b10a096b733795fdce" }
+p3-commit = { git = "https://github.com/Plonky3/Plonky3", rev = "1e87146ebfaedc2150b635b10a096b733795fdce" }
 serial_test = "3.0.0"
 
 [features]

wrappers/rust/icicle-fields/icicle-babybear/src/lib.rs

@@ -2,4 +2,5 @@ pub mod field;
 pub mod ntt;
 pub mod polynomials;
 pub mod poseidon2;
+pub mod tree;
 pub mod vec_ops;

wrappers/rust/icicle-fields/icicle-babybear/src/poseidon2/mod.rs

@@ -1,8 +1,9 @@
 use crate::field::{ScalarCfg, ScalarField};
 
 use icicle_core::error::IcicleResult;
+use icicle_core::hash::SpongeConfig;
 use icicle_core::impl_poseidon2;
-use icicle_core::poseidon2::{DiffusionStrategy, MdsType, Poseidon2, Poseidon2Config, Poseidon2Constants};
+use icicle_core::poseidon2::{DiffusionStrategy, MdsType, Poseidon2Handle, Poseidon2Impl};
 use icicle_core::traits::IcicleResultWrap;
 use icicle_cuda_runtime::device_context::DeviceContext;
 use icicle_cuda_runtime::error::CudaError;
@@ -16,14 +17,14 @@ impl_poseidon2!("babybear", babybear, ScalarField, ScalarCfg);
 pub(crate) mod tests {
     use crate::field::ScalarField;
     use icicle_core::impl_poseidon2_tests;
-    use icicle_core::poseidon2::{create_poseidon2_constants, tests::*, DiffusionStrategy, MdsType};
+    use icicle_core::poseidon2::{tests::*, DiffusionStrategy, MdsType, Poseidon2};
     use icicle_core::traits::FieldImpl;
     use icicle_cuda_runtime::device_context::DeviceContext;
 
     use p3_baby_bear::BabyBear;
     use p3_baby_bear::DiffusionMatrixBabyBear;
     use p3_field::{AbstractField, PrimeField32};
-    use p3_poseidon2::{Poseidon2, Poseidon2ExternalMatrixGeneral};
+    use p3_poseidon2::{Poseidon2 as PlonkyPoseidon2, Poseidon2ExternalMatrixGeneral};
     use p3_symmetric::Permutation;
 
     impl_poseidon2_tests!(ScalarField);
@@ -57,12 +58,13 @@ pub(crate) mod tests {
             ScalarField::from_hex("0x57a99864"),
         ];
 
-        let constants = init_poseidon::<ScalarField>(24, MdsType::Default, DiffusionStrategy::Default);
-        check_poseidon_kats(24, &kats, &constants);
+        let poseidon = init_poseidon::<ScalarField>(24, MdsType::Default, DiffusionStrategy::Default);
+        check_poseidon_kats(24, &kats, &poseidon);
     }
 
-    #[test]
-    fn test_poseidon2_plonky3_t16() {
+    type PlonkyPoseidon2T16 = PlonkyPoseidon2<BabyBear, Poseidon2ExternalMatrixGeneral, DiffusionMatrixBabyBear, 16, 7>;
+
+    pub(crate) fn get_plonky3_poseidon2_t16(rate: usize) -> (Poseidon2<ScalarField>, PlonkyPoseidon2T16) {
         let rounds_p = 13;
         let rounds_f = 8;
         const ALPHA: u64 = 7;
@@ -232,27 +234,20 @@ pub(crate) mod tests {
             cnv(605745517),
         ];
 
-        let poseidon2: Poseidon2<BabyBear, Poseidon2ExternalMatrixGeneral, DiffusionMatrixBabyBear, WIDTH, ALPHA> =
-            Poseidon2::new(
-                rounds_f,
-                external_constants.clone(),
-                Poseidon2ExternalMatrixGeneral::default(),
-                rounds_p,
-                internal_constants.clone(),
-                DiffusionMatrixBabyBear::default(),
-            );
-
-        let mut input: [BabyBear; WIDTH] = [BabyBear::zero(); WIDTH];
-        for i in 0..WIDTH {
-            input[i] = BabyBear::from_canonical_u32(i as u32);
-        }
-
-        let output = poseidon2.permute(input);
-
-        let mut kats: [ScalarField; WIDTH] = [ScalarField::zero(); WIDTH];
-        for i in 0..WIDTH {
-            kats[i] = ScalarField::from_u32(output[i].as_canonical_u32());
-        }
+        let plonky_poseidon2: PlonkyPoseidon2<
+            BabyBear,
+            Poseidon2ExternalMatrixGeneral,
+            DiffusionMatrixBabyBear,
+            WIDTH,
+            ALPHA,
+        > = PlonkyPoseidon2::new(
+            rounds_f,
+            external_constants.clone(),
+            Poseidon2ExternalMatrixGeneral::default(),
+            rounds_p,
+            internal_constants.clone(),
+            DiffusionMatrixBabyBear::default(),
+        );
 
         let ctx = DeviceContext::default();
         let mut round_constants = vec![ScalarField::zero(); rounds_f * WIDTH + rounds_p];
@@ -291,19 +286,43 @@ pub(crate) mod tests {
             ScalarField::from_u32(1 << 13),
             ScalarField::from_u32(1 << 15),
         ];
-        let constants = create_poseidon2_constants(
-            WIDTH as u32,
+
+        let poseidon = Poseidon2::new(
+            WIDTH,
+            rate,
             ALPHA as u32,
-            &ctx,
             rounds_p as u32,
             rounds_f as u32,
             &mut round_constants,
             &mut internal_matrix_diag,
             MdsType::Plonky,
             DiffusionStrategy::Montgomery,
+            &ctx,
         )
         .unwrap();
-        check_poseidon_kats(WIDTH, &kats, &constants);
+
+        (poseidon, plonky_poseidon2)
+    }
+
+    #[test]
+    fn test_poseidon2_plonky3_t16() {
+        const WIDTH: usize = 16;
+
+        let (poseidon, plonky_poseidon2) = get_plonky3_poseidon2_t16(16);
+
+        let mut input: [BabyBear; WIDTH] = [BabyBear::zero(); WIDTH];
+        for i in 0..WIDTH {
+            input[i] = BabyBear::from_canonical_u32(i as u32);
+        }
+
+        let output = plonky_poseidon2.permute(input);
+
+        let mut kats: [ScalarField; WIDTH] = [ScalarField::zero(); WIDTH];
+        for i in 0..WIDTH {
+            kats[i] = ScalarField::from_u32(output[i].as_canonical_u32());
+        }
+
+        check_poseidon_kats(WIDTH, &kats, &poseidon);
     }
 
     #[test]
@@ -549,22 +568,27 @@ pub(crate) mod tests {
             cnv(1810596765),
         ];
 
-        let poseidon2: Poseidon2<BabyBear, Poseidon2ExternalMatrixGeneral, DiffusionMatrixBabyBear, WIDTH, ALPHA> =
-            Poseidon2::new(
-                rounds_f,
-                external_constants.clone(),
-                Poseidon2ExternalMatrixGeneral::default(),
-                rounds_p,
-                internal_constants.clone(),
-                DiffusionMatrixBabyBear::default(),
-            );
+        let plonky_poseidon2: PlonkyPoseidon2<
+            BabyBear,
+            Poseidon2ExternalMatrixGeneral,
+            DiffusionMatrixBabyBear,
+            WIDTH,
+            ALPHA,
+        > = PlonkyPoseidon2::new(
+            rounds_f,
+            external_constants.clone(),
+            Poseidon2ExternalMatrixGeneral::default(),
+            rounds_p,
+            internal_constants.clone(),
+            DiffusionMatrixBabyBear::default(),
+        );
 
         let mut input: [BabyBear; WIDTH] = [BabyBear::zero(); WIDTH];
         for i in 0..WIDTH {
             input[i] = BabyBear::from_canonical_u32(i as u32);
         }
 
-        let output = poseidon2.permute(input);
+        let output = plonky_poseidon2.permute(input);
 
         let mut kats: [ScalarField; WIDTH] = [ScalarField::zero(); WIDTH];
         for i in 0..WIDTH {
@@ -616,18 +640,19 @@ pub(crate) mod tests {
             ScalarField::from_u32(1 << 22),
             ScalarField::from_u32(1 << 23),
         ];
-        let constants = create_poseidon2_constants(
-            WIDTH as u32,
+        let poseidon = Poseidon2::new(
+            WIDTH,
+            24,
             ALPHA as u32,
-            &ctx,
             rounds_p as u32,
             rounds_f as u32,
             &mut round_constants,
             &mut internal_matrix_diag,
             MdsType::Plonky,
             DiffusionStrategy::Montgomery,
+            &ctx,
         )
         .unwrap();
-        check_poseidon_kats(WIDTH, &kats, &constants);
+        check_poseidon_kats(WIDTH, &kats, &poseidon);
     }
 }

wrappers/rust/icicle-fields/icicle-babybear/src/tree/mmcs.rs

@@ -0,0 +1,125 @@
+use icicle_core::error::IcicleResult;
+use icicle_core::hash::SpongeHash;
+use icicle_core::traits::IcicleResultWrap;
+use icicle_core::tree::{mmcs::FieldMmcs, TreeBuilderConfig};
+use icicle_core::{impl_mmcs, Matrix};
+use icicle_cuda_runtime::memory::{HostOrDeviceSlice, HostSlice};
+use std::ffi::c_void;
+
+use crate::field::ScalarField;
+
+impl_mmcs!("babybear", babybear_mmcs, ScalarField, ScalarCfg, BabyBearMmcs);
+
+#[cfg(test)]
+pub(crate) mod tests {
+    use std::ffi::c_void;
+
+    use icicle_core::{
+        ntt::FieldImpl,
+        tree::{merkle_tree_digests_len, TreeBuilderConfig},
+        Matrix,
+    };
+    use icicle_cuda_runtime::memory::HostSlice;
+    use p3_baby_bear::{BabyBear, DiffusionMatrixBabyBear};
+    use p3_commit::Mmcs;
+    use p3_field::{AbstractField, Field};
+    use p3_matrix::dense::RowMajorMatrix;
+    use p3_merkle_tree::FieldMerkleTreeMmcs;
+    use p3_poseidon2::{Poseidon2 as PlonkyPoseidon2, Poseidon2ExternalMatrixGeneral};
+    use p3_symmetric::{PaddingFreeSponge, TruncatedPermutation};
+
+    use crate::{
+        field::ScalarField,
+        poseidon2::tests::get_plonky3_poseidon2_t16,
+        tree::mmcs::{BabyBearMmcs, FieldMmcs},
+    };
+
+    type PlonkyPoseidon2T16 = PlonkyPoseidon2<BabyBear, Poseidon2ExternalMatrixGeneral, DiffusionMatrixBabyBear, 16, 7>;
+
+    #[test]
+    fn test_poseidon2_mmcs_plonky3() {
+        const WIDTH: usize = 16;
+        const RATE: usize = 8;
+        const ARITY: usize = 2;
+        const HEIGHT: usize = 15;
+        const ROWS: usize = 1 << HEIGHT;
+        const COLS: usize = 32;
+        const DIGEST_ELEMENTS: usize = 8;
+
+        let (poseidon, plonky_poseidon2) = get_plonky3_poseidon2_t16(RATE);
+
+        type H = PaddingFreeSponge<PlonkyPoseidon2T16, WIDTH, RATE, RATE>;
+        let h = H::new(plonky_poseidon2.clone());
+
+        type C = TruncatedPermutation<PlonkyPoseidon2T16, ARITY, RATE, WIDTH>;
+        let c = C::new(plonky_poseidon2.clone());
+
+        type F = BabyBear;
+
+        let mut input = vec![F::zero(); ROWS * COLS];
+        let mut icicle_input = vec![ScalarField::zero(); ROWS * COLS];
+        for i in 0..ROWS * COLS {
+            input[i] = F::from_canonical_u32(i as u32);
+            icicle_input[i] = ScalarField::from_u32(i as u32);
+        }
+
+        let mut input2 = vec![F::zero(); (ROWS / 2) * COLS];
+        let mut icicle_input2 = vec![ScalarField::zero(); (ROWS / 2) * COLS];
+        for i in 0..(ROWS / 2) * COLS {
+            input2[i] = F::from_canonical_u32(i as u32);
+            icicle_input2[i] = ScalarField::from_u32(i as u32);
+        }
+
+        let matrix = RowMajorMatrix::new(input.clone(), COLS);
+        let matrix2 = RowMajorMatrix::new(input2.clone(), COLS);
+        let leaves = vec![matrix, matrix2];
+        // let leaves = vec![matrix];
+
+        let mmcs =
+            FieldMerkleTreeMmcs::<<F as Field>::Packing, <F as Field>::Packing, H, C, DIGEST_ELEMENTS>::new(h, c);
+
+        let (commit, _data) = mmcs.commit(leaves);
+
+        let mut config = TreeBuilderConfig::default();
+        config.arity = ARITY as u32;
+        config.keep_rows = HEIGHT as u32 + 1;
+        config.digest_elements = DIGEST_ELEMENTS as u32;
+        let digests_len = merkle_tree_digests_len(HEIGHT as u32, ARITY as u32, DIGEST_ELEMENTS as u32);
+        let mut digests = vec![ScalarField::zero(); digests_len];
+        // let mut digests = vec![ScalarField::zero(); COLS];
+
+        let leaves_slice = vec![
+            Matrix {
+                values: icicle_input.as_ptr() as *const c_void,
+                width: COLS,
+                height: ROWS,
+            },
+            Matrix {
+                values: icicle_input2.as_ptr() as *const c_void,
+                width: COLS,
+                height: ROWS / 2,
+            },
+        ];
+        let digests_slice = HostSlice::from_mut_slice(&mut digests);
+
+        BabyBearMmcs::mmcs_commit(leaves_slice, digests_slice, &poseidon, &poseidon, &config).unwrap();
+
+        let mut converted = vec![BabyBear::zero(); digests_len];
+        for i in 0..digests_len {
+            let mut scalar_bytes = [0u8; 4];
+            scalar_bytes.copy_from_slice(&digests_slice[i].to_bytes_le());
+            converted[i] = BabyBear::from_canonical_u32(u32::from_le_bytes(scalar_bytes));
+        }
+
+        // println!("Plonky: {:?}", _data);
+        // println!("Icicle: {:?}", converted);
+        // assert_eq!(commit, converted);
+
+        let commit_vec: Vec<BabyBear> = commit
+            .into_iter()
+            .collect();
+        for i in 0..DIGEST_ELEMENTS {
+            assert_eq!(converted[converted.len() - DIGEST_ELEMENTS + i], commit_vec[i]);
+        }
+    }
+}

wrappers/rust/icicle-fields/icicle-babybear/src/tree/mod.rs

@@ -0,0 +1,107 @@
+use icicle_core::error::IcicleResult;
+use icicle_core::hash::SpongeHash;
+use icicle_core::impl_field_tree_builder;
+use icicle_core::traits::IcicleResultWrap;
+use icicle_core::tree::{FieldTreeBuilder, TreeBuilderConfig};
+use icicle_cuda_runtime::memory::HostOrDeviceSlice;
+use std::ffi::c_void;
+
+use crate::field::ScalarField;
+
+pub mod mmcs;
+
+impl_field_tree_builder!("babybear", babybear_tb, ScalarField, ScalarCfg, BabyBearTreeBuilder);
+
+#[cfg(test)]
+pub(crate) mod tests {
+    use icicle_core::{
+        ntt::FieldImpl,
+        poseidon2::{DiffusionStrategy, MdsType, Poseidon2},
+        tree::{tests::check_build_field_merkle_tree, FieldTreeBuilder, TreeBuilderConfig},
+    };
+    use icicle_cuda_runtime::device_context;
+    use icicle_cuda_runtime::memory::HostSlice;
+    use p3_baby_bear::{BabyBear, DiffusionMatrixBabyBear};
+    use p3_commit::Mmcs;
+    use p3_field::{AbstractField, Field};
+    use p3_matrix::dense::RowMajorMatrix;
+    use p3_merkle_tree::FieldMerkleTreeMmcs;
+    use p3_poseidon2::{Poseidon2 as PlonkyPoseidon2, Poseidon2ExternalMatrixGeneral};
+    use p3_symmetric::{PaddingFreeSponge, TruncatedPermutation};
+
+    use crate::{field::ScalarField, poseidon2::tests::get_plonky3_poseidon2_t16, tree::BabyBearTreeBuilder};
+
+    #[test]
+    fn poseidon2_merkle_tree_test() {
+        let ctx = device_context::DeviceContext::default();
+        let sponge = Poseidon2::load(2, 2, MdsType::Default, DiffusionStrategy::Default, &ctx).unwrap();
+
+        check_build_field_merkle_tree::<_, _, BabyBearTreeBuilder>(25, 2, &sponge, &sponge, ScalarField::zero());
+    }
+
+    type PlonkyPoseidon2T16 = PlonkyPoseidon2<BabyBear, Poseidon2ExternalMatrixGeneral, DiffusionMatrixBabyBear, 16, 7>;
+
+    #[test]
+    fn test_poseidon2_tree_plonky3() {
+        const WIDTH: usize = 16;
+        const ARITY: usize = 2;
+        const HEIGHT: usize = 15;
+        const ROWS: usize = 1 << HEIGHT;
+        const COLS: usize = 8;
+
+        let (poseidon, plonky_poseidon2) = get_plonky3_poseidon2_t16(8);
+
+        type H = PaddingFreeSponge<PlonkyPoseidon2T16, WIDTH, COLS, COLS>;
+        let h = H::new(plonky_poseidon2.clone());
+
+        type C = TruncatedPermutation<PlonkyPoseidon2T16, ARITY, COLS, WIDTH>;
+        let c = C::new(plonky_poseidon2.clone());
+
+        type F = BabyBear;
+
+        let mut input = vec![F::zero(); ROWS * COLS];
+        let mut icicle_input = vec![ScalarField::zero(); ROWS * COLS];
+        for i in 0..ROWS * COLS {
+            input[i] = F::from_canonical_u32(i as u32);
+            icicle_input[i] = ScalarField::from_u32(i as u32);
+        }
+
+        let matrix = RowMajorMatrix::new(input, COLS);
+        let leaves = vec![matrix];
+
+        let mmcs = FieldMerkleTreeMmcs::<<F as Field>::Packing, <F as Field>::Packing, H, C, 8>::new(h, c);
+
+        let (commit, _data) = mmcs.commit(leaves);
+
+        let mut config = TreeBuilderConfig::default();
+        config.arity = ARITY as u32;
+        config.keep_rows = 1;
+        config.digest_elements = COLS as u32;
+        let input_block_len = COLS;
+        // let digests_len = merkle_tree_digests_len(2 as u32, ARITY as u32, COLS as u32);
+        // let mut digests = vec![ScalarField::zero(); digests_len];
+        let mut digests = vec![ScalarField::zero(); COLS];
+
+        let leaves_slice = HostSlice::from_slice(&icicle_input);
+        let digests_slice = HostSlice::from_mut_slice(&mut digests);
+
+        BabyBearTreeBuilder::build_merkle_tree(
+            leaves_slice,
+            digests_slice,
+            HEIGHT,
+            input_block_len,
+            &poseidon,
+            &poseidon,
+            &config,
+        )
+        .unwrap();
+
+        let mut converted: [BabyBear; COLS] = [BabyBear::zero(); COLS];
+        for i in 0..COLS {
+            let mut scalar_bytes = [0u8; 4];
+            scalar_bytes.copy_from_slice(&digests_slice[i].to_bytes_le());
+            converted[i] = BabyBear::from_canonical_u32(u32::from_le_bytes(scalar_bytes));
+        }
+        assert_eq!(commit, converted);
+    }
+}