mirror of
https://github.com/luxfi/crypto.git
synced 2026-07-27 01:54:50 +00:00
feat: add GPU stub files for CGO-less builds
Add stub implementations for GPU packages that return false for Available() when GPU libraries are not present. This allows building with CGO_ENABLED=0 or CGO_ENABLED=1 without requiring Metal/CUDA libraries. - bls/gpu/gpu_stub.go - hash/blake3/gpu/gpu_stub.go - hash/poseidon2/gpu/gpu_stub.go - pqcrypto/*/gpu/gpu_stub.go (mldsa, mlkem, slhdsa)
This commit is contained in:
Zach Kelling
@@ -0,0 +1,214 @@
|
||||
// Copyright (C) 2025, Lux Industries, Inc. All rights reserved.
|
||||
// See the file LICENSE for licensing terms.
|
||||
|
||||
//go:build cgo && darwin
|
||||
// +build cgo,darwin
|
||||
|
||||
// Package blake3 provides Blake3 hash functions via luxcpp/crypto.
|
||||
// This CGO version provides optimized hashing with automatic CPU/GPU selection.
|
||||
// For GPU-only operations, use the blake3/gpu subpackage.
|
||||
package blake3
|
||||
|
||||
/*
|
||||
#cgo CFLAGS: -I${SRCDIR}/../../../../luxcpp/crypto/include
|
||||
#cgo LDFLAGS: -L${SRCDIR}/../../../../luxcpp/crypto/build-local -lluxcrypto -framework Metal -framework Foundation
|
||||
|
||||
#include <stdint.h>
|
||||
#include <stdlib.h>
|
||||
|
||||
// C API from luxcpp/crypto
|
||||
extern void crypto_blake3(uint8_t* out, const uint8_t* in, size_t len);
|
||||
extern int crypto_batch_hash(uint8_t** outs, const uint8_t* const* ins, const size_t* lens, uint32_t count, int hash_type);
|
||||
extern int crypto_gpu_available(void);
|
||||
extern const char* crypto_get_backend(void);
|
||||
*/
|
||||
import "C"
|
||||
|
||||
import (
|
||||
"encoding/binary"
|
||||
"unsafe"
|
||||
)
|
||||
|
||||
const (
|
||||
// Size256 is the standard Blake3 output size
|
||||
Size256 = 32
|
||||
// Size512 is extended Blake3 output size
|
||||
Size512 = 64
|
||||
// gpuThreshold is minimum batch size for GPU acceleration
|
||||
gpuThreshold = 64
|
||||
)
|
||||
|
||||
var (
|
||||
// useGPU indicates if GPU acceleration is available
|
||||
useGPU = C.crypto_gpu_available() != 0
|
||||
)
|
||||
|
||||
// GPUAvailable returns true if GPU acceleration is available.
|
||||
func GPUAvailable() bool {
|
||||
return useGPU
|
||||
}
|
||||
|
||||
// Backend returns the active backend name ("Metal", "CUDA", or "CPU").
|
||||
func Backend() string {
|
||||
return C.GoString(C.crypto_get_backend())
|
||||
}
|
||||
|
||||
// Sum256 computes a 32-byte Blake3 hash.
|
||||
func Sum256(data []byte) [Size256]byte {
|
||||
var out [Size256]byte
|
||||
if len(data) == 0 {
|
||||
C.crypto_blake3((*C.uint8_t)(&out[0]), nil, 0)
|
||||
return out
|
||||
}
|
||||
C.crypto_blake3(
|
||||
(*C.uint8_t)(&out[0]),
|
||||
(*C.uint8_t)(&data[0]),
|
||||
C.size_t(len(data)),
|
||||
)
|
||||
return out
|
||||
}
|
||||
|
||||
// Sum512 computes a 64-byte Blake3 hash using XOF mode.
|
||||
// Implemented as XOF by hashing with counter blocks.
|
||||
func Sum512(data []byte) [Size512]byte {
|
||||
var out [Size512]byte
|
||||
xof := sumXOF(data, Size512)
|
||||
copy(out[:], xof)
|
||||
return out
|
||||
}
|
||||
|
||||
// SumXOF computes an extendable output hash of arbitrary length.
|
||||
func SumXOF(data []byte, outLen int) []byte {
|
||||
if outLen <= 0 {
|
||||
return nil
|
||||
}
|
||||
return sumXOF(data, outLen)
|
||||
}
|
||||
|
||||
// sumXOF implements XOF by chaining hash outputs with counter.
|
||||
func sumXOF(data []byte, outLen int) []byte {
|
||||
if outLen <= Size256 {
|
||||
hash := Sum256(data)
|
||||
return hash[:outLen]
|
||||
}
|
||||
|
||||
out := make([]byte, outLen)
|
||||
// First block: hash the data directly
|
||||
first := Sum256(data)
|
||||
copy(out[:Size256], first[:])
|
||||
|
||||
// Subsequent blocks: hash(data || counter)
|
||||
remaining := outLen - Size256
|
||||
counter := uint64(1)
|
||||
offset := Size256
|
||||
|
||||
for remaining > 0 {
|
||||
// Create input: data || counter (little-endian)
|
||||
counterBytes := make([]byte, 8)
|
||||
binary.LittleEndian.PutUint64(counterBytes, counter)
|
||||
input := append(data, counterBytes...)
|
||||
|
||||
block := Sum256(input)
|
||||
toCopy := Size256
|
||||
if remaining < Size256 {
|
||||
toCopy = remaining
|
||||
}
|
||||
copy(out[offset:offset+toCopy], block[:toCopy])
|
||||
|
||||
offset += toCopy
|
||||
remaining -= toCopy
|
||||
counter++
|
||||
}
|
||||
|
||||
return out
|
||||
}
|
||||
|
||||
// BatchSum256 computes Blake3 hashes for multiple inputs in parallel.
|
||||
// More efficient than calling Sum256 repeatedly for batch operations.
|
||||
func BatchSum256(inputs [][]byte) [][Size256]byte {
|
||||
n := len(inputs)
|
||||
if n == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
// For small batches, use sequential processing
|
||||
if n < gpuThreshold || !useGPU {
|
||||
results := make([][Size256]byte, n)
|
||||
for i, data := range inputs {
|
||||
results[i] = Sum256(data)
|
||||
}
|
||||
return results
|
||||
}
|
||||
|
||||
// Prepare C arrays for batch processing
|
||||
outs := make([]*C.uint8_t, n)
|
||||
ins := make([]*C.uint8_t, n)
|
||||
lens := make([]C.size_t, n)
|
||||
|
||||
results := make([][Size256]byte, n)
|
||||
for i, data := range inputs {
|
||||
outs[i] = (*C.uint8_t)(&results[i][0])
|
||||
if len(data) > 0 {
|
||||
ins[i] = (*C.uint8_t)(&data[0])
|
||||
}
|
||||
lens[i] = C.size_t(len(data))
|
||||
}
|
||||
|
||||
// Call batch hash (hash_type 2 = BLAKE3)
|
||||
C.crypto_batch_hash(
|
||||
(**C.uint8_t)(unsafe.Pointer(&outs[0])),
|
||||
(**C.uint8_t)(unsafe.Pointer(&ins[0])),
|
||||
(*C.size_t)(&lens[0]),
|
||||
C.uint32_t(n),
|
||||
2, // BLAKE3
|
||||
)
|
||||
|
||||
return results
|
||||
}
|
||||
|
||||
// MerkleRoot computes the Merkle root of leaves using Blake3.
|
||||
// This is a software implementation; for GPU-accelerated version, use blake3/gpu.
|
||||
func MerkleRoot(leaves [][]byte) [Size256]byte {
|
||||
var out [Size256]byte
|
||||
n := len(leaves)
|
||||
if n == 0 {
|
||||
return out
|
||||
}
|
||||
|
||||
// Hash each leaf first
|
||||
hashes := make([][Size256]byte, n)
|
||||
for i, leaf := range leaves {
|
||||
hashes[i] = Sum256(leaf)
|
||||
}
|
||||
|
||||
// Build tree bottom-up
|
||||
for len(hashes) > 1 {
|
||||
nextLevel := make([][Size256]byte, (len(hashes)+1)/2)
|
||||
for i := 0; i < len(hashes); i += 2 {
|
||||
if i+1 < len(hashes) {
|
||||
// Hash pair
|
||||
combined := make([]byte, Size256*2)
|
||||
copy(combined[:Size256], hashes[i][:])
|
||||
copy(combined[Size256:], hashes[i+1][:])
|
||||
nextLevel[i/2] = Sum256(combined)
|
||||
} else {
|
||||
// Odd node, promote to next level
|
||||
nextLevel[i/2] = hashes[i]
|
||||
}
|
||||
}
|
||||
hashes = nextLevel
|
||||
}
|
||||
|
||||
if len(hashes) > 0 {
|
||||
out = hashes[0]
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// DeriveKey derives a key from context and key material using Blake3 KDF.
|
||||
func DeriveKey(context string, keyMaterial []byte) [Size256]byte {
|
||||
// Combine context and key material with domain separation
|
||||
input := append([]byte(context), 0x00)
|
||||
input = append(input, keyMaterial...)
|
||||
return Sum256(input)
|
||||
}
|
||||
@@ -0,0 +1,177 @@
|
||||
// Copyright (C) 2025, Lux Industries, Inc. All rights reserved.
|
||||
// See the file LICENSE for licensing terms.
|
||||
|
||||
//go:build luxgpu && darwin
|
||||
|
||||
// Package gpu provides GPU-accelerated Blake3 hash functions via Metal.
|
||||
// This package links to luxcpp/crypto for hardware acceleration.
|
||||
// The C++ library handles automatic fallback to CPU when GPU is unavailable.
|
||||
package gpu
|
||||
|
||||
/*
|
||||
#cgo CFLAGS: -I${SRCDIR}/../../../../../luxcpp/crypto/include
|
||||
#cgo LDFLAGS: -L${SRCDIR}/../../../../../luxcpp/crypto/build-local -lluxcrypto -framework Metal -framework Foundation
|
||||
|
||||
#include <stdint.h>
|
||||
#include <stdlib.h>
|
||||
|
||||
// C API from luxcpp/crypto
|
||||
extern void crypto_blake3(uint8_t* out, const uint8_t* in, size_t len);
|
||||
extern int crypto_batch_hash(uint8_t** outs, const uint8_t* const* ins, const size_t* lens, uint32_t count, int hash_type);
|
||||
extern int crypto_gpu_available(void);
|
||||
extern const char* crypto_get_backend(void);
|
||||
|
||||
// Metal-specific Blake3 functions
|
||||
extern int metal_blake3_hash256(const uint8_t* data, uint32_t len, uint8_t* out);
|
||||
extern int metal_blake3_hash512(const uint8_t* data, uint32_t len, uint8_t* out);
|
||||
extern int metal_blake3_xof(const uint8_t* data, uint32_t inLen, uint8_t* out, uint32_t outLen);
|
||||
extern int metal_blake3_merkle_root(const uint8_t* leaves, uint32_t numLeaves, uint8_t* out);
|
||||
extern int metal_blake3_derive_key(const uint8_t* context, uint32_t contextLen, const uint8_t* key, uint32_t keyLen, uint8_t* out);
|
||||
*/
|
||||
import "C"
|
||||
|
||||
import (
|
||||
"unsafe"
|
||||
)
|
||||
|
||||
const (
|
||||
// Size256 is the standard Blake3 output size
|
||||
Size256 = 32
|
||||
// Size512 is extended Blake3 output size
|
||||
Size512 = 64
|
||||
// BatchThreshold is minimum batch size for GPU acceleration
|
||||
BatchThreshold = 64
|
||||
)
|
||||
|
||||
// Available returns true if GPU acceleration is available.
|
||||
func Available() bool {
|
||||
return C.crypto_gpu_available() != 0
|
||||
}
|
||||
|
||||
// Backend returns the active backend name ("Metal", "CUDA", or "CPU").
|
||||
func Backend() string {
|
||||
return C.GoString(C.crypto_get_backend())
|
||||
}
|
||||
|
||||
// Sum256 computes a 32-byte Blake3 hash using GPU if available.
|
||||
func Sum256(data []byte) [Size256]byte {
|
||||
var out [Size256]byte
|
||||
if len(data) == 0 {
|
||||
C.crypto_blake3((*C.uint8_t)(&out[0]), nil, 0)
|
||||
return out
|
||||
}
|
||||
C.crypto_blake3(
|
||||
(*C.uint8_t)(&out[0]),
|
||||
(*C.uint8_t)(&data[0]),
|
||||
C.size_t(len(data)),
|
||||
)
|
||||
return out
|
||||
}
|
||||
|
||||
// Sum512 computes a 64-byte Blake3 hash using Metal GPU.
|
||||
func Sum512(data []byte) [Size512]byte {
|
||||
var out [Size512]byte
|
||||
var dataPtr *C.uint8_t
|
||||
if len(data) > 0 {
|
||||
dataPtr = (*C.uint8_t)(&data[0])
|
||||
}
|
||||
C.metal_blake3_hash512(dataPtr, C.uint32_t(len(data)), (*C.uint8_t)(&out[0]))
|
||||
return out
|
||||
}
|
||||
|
||||
// SumXOF computes an extendable output hash of arbitrary length.
|
||||
func SumXOF(data []byte, outLen int) []byte {
|
||||
if outLen <= 0 {
|
||||
return nil
|
||||
}
|
||||
out := make([]byte, outLen)
|
||||
var dataPtr *C.uint8_t
|
||||
if len(data) > 0 {
|
||||
dataPtr = (*C.uint8_t)(&data[0])
|
||||
}
|
||||
C.metal_blake3_xof(
|
||||
dataPtr,
|
||||
C.uint32_t(len(data)),
|
||||
(*C.uint8_t)(&out[0]),
|
||||
C.uint32_t(outLen),
|
||||
)
|
||||
return out
|
||||
}
|
||||
|
||||
// BatchSum256 computes Blake3 hashes for multiple inputs in parallel on GPU.
|
||||
// More efficient than calling Sum256 repeatedly for batch operations.
|
||||
func BatchSum256(inputs [][]byte) [][Size256]byte {
|
||||
n := len(inputs)
|
||||
if n == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
// Prepare C arrays for batch processing
|
||||
outs := make([]*C.uint8_t, n)
|
||||
ins := make([]*C.uint8_t, n)
|
||||
lens := make([]C.size_t, n)
|
||||
|
||||
results := make([][Size256]byte, n)
|
||||
for i, data := range inputs {
|
||||
outs[i] = (*C.uint8_t)(&results[i][0])
|
||||
if len(data) > 0 {
|
||||
ins[i] = (*C.uint8_t)(&data[0])
|
||||
}
|
||||
lens[i] = C.size_t(len(data))
|
||||
}
|
||||
|
||||
// Call batch hash (hash_type 2 = BLAKE3)
|
||||
C.crypto_batch_hash(
|
||||
(**C.uint8_t)(unsafe.Pointer(&outs[0])),
|
||||
(**C.uint8_t)(unsafe.Pointer(&ins[0])),
|
||||
(*C.size_t)(&lens[0]),
|
||||
C.uint32_t(n),
|
||||
2, // BLAKE3
|
||||
)
|
||||
|
||||
return results
|
||||
}
|
||||
|
||||
// MerkleRoot computes the Merkle root of leaves using GPU-accelerated Blake3.
|
||||
func MerkleRoot(leaves [][]byte) [Size256]byte {
|
||||
var out [Size256]byte
|
||||
n := len(leaves)
|
||||
if n == 0 {
|
||||
return out
|
||||
}
|
||||
|
||||
// Flatten leaves to contiguous buffer
|
||||
leafData := make([]byte, n*Size256)
|
||||
for i, leaf := range leaves {
|
||||
hash := Sum256(leaf)
|
||||
copy(leafData[i*Size256:], hash[:])
|
||||
}
|
||||
|
||||
C.metal_blake3_merkle_root(
|
||||
(*C.uint8_t)(&leafData[0]),
|
||||
C.uint32_t(n),
|
||||
(*C.uint8_t)(&out[0]),
|
||||
)
|
||||
return out
|
||||
}
|
||||
|
||||
// DeriveKey derives a key from context and key material using Blake3 KDF.
|
||||
func DeriveKey(context string, keyMaterial []byte) [Size256]byte {
|
||||
var out [Size256]byte
|
||||
contextBytes := []byte(context)
|
||||
|
||||
var contextPtr, keyPtr *C.uint8_t
|
||||
if len(contextBytes) > 0 {
|
||||
contextPtr = (*C.uint8_t)(&contextBytes[0])
|
||||
}
|
||||
if len(keyMaterial) > 0 {
|
||||
keyPtr = (*C.uint8_t)(&keyMaterial[0])
|
||||
}
|
||||
|
||||
C.metal_blake3_derive_key(
|
||||
contextPtr, C.uint32_t(len(contextBytes)),
|
||||
keyPtr, C.uint32_t(len(keyMaterial)),
|
||||
(*C.uint8_t)(&out[0]),
|
||||
)
|
||||
return out
|
||||
}
|
||||
@@ -0,0 +1,97 @@
|
||||
// Copyright (C) 2025, Lux Industries, Inc. All rights reserved.
|
||||
// See the file LICENSE for licensing terms.
|
||||
|
||||
//go:build !luxgpu
|
||||
|
||||
// Package gpu provides GPU-accelerated Blake3 hash functions.
|
||||
// This is the stub implementation when GPU acceleration is not available.
|
||||
package gpu
|
||||
|
||||
import (
|
||||
"lukechampine.com/blake3"
|
||||
)
|
||||
|
||||
const (
|
||||
// Size256 is the standard Blake3 output size
|
||||
Size256 = 32
|
||||
// Size512 is extended Blake3 output size
|
||||
Size512 = 64
|
||||
// BatchThreshold is minimum batch size for GPU acceleration
|
||||
BatchThreshold = 64
|
||||
)
|
||||
|
||||
// Available returns true if GPU acceleration is available.
|
||||
func Available() bool {
|
||||
return false
|
||||
}
|
||||
|
||||
// Backend returns the active backend name.
|
||||
func Backend() string {
|
||||
return "CPU"
|
||||
}
|
||||
|
||||
// Sum256 computes a 32-byte Blake3 hash.
|
||||
func Sum256(data []byte) [Size256]byte {
|
||||
return blake3.Sum256(data)
|
||||
}
|
||||
|
||||
// Sum512 computes a 64-byte Blake3 hash.
|
||||
func Sum512(data []byte) [Size512]byte {
|
||||
return blake3.Sum512(data)
|
||||
}
|
||||
|
||||
// SumXOF computes an extendable output hash of arbitrary length.
|
||||
func SumXOF(data []byte, outLen int) []byte {
|
||||
if outLen <= 0 {
|
||||
return nil
|
||||
}
|
||||
h := blake3.New(outLen, nil)
|
||||
h.Write(data)
|
||||
return h.Sum(nil)
|
||||
}
|
||||
|
||||
// BatchSum256 computes multiple 32-byte hashes.
|
||||
func BatchSum256(inputs [][]byte) [][Size256]byte {
|
||||
results := make([][Size256]byte, len(inputs))
|
||||
for i, input := range inputs {
|
||||
results[i] = blake3.Sum256(input)
|
||||
}
|
||||
return results
|
||||
}
|
||||
|
||||
// MerkleRoot computes the Merkle root of leaves using Blake3.
|
||||
func MerkleRoot(leaves [][]byte) [Size256]byte {
|
||||
if len(leaves) == 0 {
|
||||
return [Size256]byte{}
|
||||
}
|
||||
if len(leaves) == 1 {
|
||||
return blake3.Sum256(leaves[0])
|
||||
}
|
||||
// Simple iterative merkle tree
|
||||
hashes := make([][Size256]byte, len(leaves))
|
||||
for i, leaf := range leaves {
|
||||
hashes[i] = blake3.Sum256(leaf)
|
||||
}
|
||||
for len(hashes) > 1 {
|
||||
next := make([][Size256]byte, (len(hashes)+1)/2)
|
||||
for i := 0; i < len(hashes); i += 2 {
|
||||
if i+1 < len(hashes) {
|
||||
combined := append(hashes[i][:], hashes[i+1][:]...)
|
||||
next[i/2] = blake3.Sum256(combined)
|
||||
} else {
|
||||
next[i/2] = hashes[i]
|
||||
}
|
||||
}
|
||||
hashes = next
|
||||
}
|
||||
return hashes[0]
|
||||
}
|
||||
|
||||
// DeriveKey derives a key using Blake3 KDF.
|
||||
func DeriveKey(context, key []byte) [Size256]byte {
|
||||
h := blake3.New(Size256, key)
|
||||
h.Write(context)
|
||||
var out [Size256]byte
|
||||
copy(out[:], h.Sum(nil))
|
||||
return out
|
||||
}
|
||||
@@ -0,0 +1,302 @@
|
||||
// Copyright (C) 2025, Lux Industries, Inc. All rights reserved.
|
||||
// See the file LICENSE for licensing terms.
|
||||
|
||||
//go:build luxgpu && darwin
|
||||
|
||||
// Package gpu provides GPU-accelerated Poseidon2 hash functions via Metal.
|
||||
// This package links to luxcpp/crypto for hardware acceleration.
|
||||
// The C++ library handles automatic fallback to CPU when GPU is unavailable.
|
||||
// Poseidon2 is a ZK-friendly hash function designed for BN254 scalar field.
|
||||
package gpu
|
||||
|
||||
/*
|
||||
#cgo CFLAGS: -I${SRCDIR}/../../../../../luxcpp/crypto/include
|
||||
#cgo LDFLAGS: -L${SRCDIR}/../../../../../luxcpp/crypto/build-local -lluxcrypto -framework Metal -framework Foundation
|
||||
|
||||
#include <stdint.h>
|
||||
#include <stdlib.h>
|
||||
#include <stdbool.h>
|
||||
#include "lux/crypto/metal_zk.h"
|
||||
*/
|
||||
import "C"
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"unsafe"
|
||||
)
|
||||
|
||||
// ElementSize is the size of a BN254 scalar field element in bytes (256-bit).
|
||||
const ElementSize = 32
|
||||
|
||||
// Element represents a 256-bit BN254 scalar field element.
|
||||
// Stored as 4 x 64-bit limbs in little-endian order.
|
||||
type Element [ElementSize]byte
|
||||
|
||||
// context is the global Metal ZK context (initialized on first use).
|
||||
var (
|
||||
ctx *C.MetalZKContext
|
||||
ctxReady bool
|
||||
)
|
||||
|
||||
// initContext lazily initializes the Metal ZK context.
|
||||
func initContext() error {
|
||||
if ctxReady {
|
||||
return nil
|
||||
}
|
||||
ctx = C.metal_zk_init()
|
||||
if ctx == nil {
|
||||
return errors.New("Metal ZK initialization failed")
|
||||
}
|
||||
ctxReady = true
|
||||
return nil
|
||||
}
|
||||
|
||||
// Available returns true if Metal GPU acceleration is available.
|
||||
func Available() bool {
|
||||
return bool(C.metal_zk_available())
|
||||
}
|
||||
|
||||
// Threshold returns the recommended batch size for GPU offload.
|
||||
func Threshold() uint32 {
|
||||
return uint32(C.metal_zk_get_threshold(C.METAL_ZK_OP_POSEIDON2_HASH))
|
||||
}
|
||||
|
||||
// HashPair computes Poseidon2 hash of two elements (2-to-1 compression).
|
||||
// This is the core operation for Merkle tree construction.
|
||||
func HashPair(left, right Element) (Element, error) {
|
||||
var out Element
|
||||
if err := initContext(); err != nil {
|
||||
return out, err
|
||||
}
|
||||
|
||||
leftFr := (*C.Fr256)(unsafe.Pointer(&left[0]))
|
||||
rightFr := (*C.Fr256)(unsafe.Pointer(&right[0]))
|
||||
outFr := (*C.Fr256)(unsafe.Pointer(&out[0]))
|
||||
|
||||
ret := C.metal_zk_poseidon2_hash_pair(ctx, outFr, leftFr, rightFr, 1)
|
||||
if ret != C.METAL_ZK_SUCCESS {
|
||||
return out, errors.New("Poseidon2 hash pair failed")
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// BatchHashPair computes Poseidon2 hashes for multiple pairs on GPU.
|
||||
// More efficient than calling HashPair repeatedly for batch operations.
|
||||
func BatchHashPair(left, right []Element) ([]Element, error) {
|
||||
n := len(left)
|
||||
if n == 0 || n != len(right) {
|
||||
return nil, errors.New("invalid input lengths")
|
||||
}
|
||||
if err := initContext(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
outputs := make([]Element, n)
|
||||
|
||||
// GPU batch processing
|
||||
leftFr := (*C.Fr256)(unsafe.Pointer(&left[0]))
|
||||
rightFr := (*C.Fr256)(unsafe.Pointer(&right[0]))
|
||||
outFr := (*C.Fr256)(unsafe.Pointer(&outputs[0]))
|
||||
|
||||
ret := C.metal_zk_poseidon2_hash_pair(ctx, outFr, leftFr, rightFr, C.uint32_t(n))
|
||||
if ret != C.METAL_ZK_SUCCESS {
|
||||
return nil, errors.New("Poseidon2 batch hash pair failed")
|
||||
}
|
||||
return outputs, nil
|
||||
}
|
||||
|
||||
// MerkleLayer computes one layer of a Merkle tree from leaf pairs.
|
||||
// Input size must be even. Output size is input size / 2.
|
||||
func MerkleLayer(current []Element) ([]Element, error) {
|
||||
n := len(current)
|
||||
if n == 0 || n%2 != 0 {
|
||||
return nil, errors.New("layer size must be even")
|
||||
}
|
||||
if err := initContext(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
output := make([]Element, n/2)
|
||||
currentFr := (*C.Fr256)(unsafe.Pointer(¤t[0]))
|
||||
outFr := (*C.Fr256)(unsafe.Pointer(&output[0]))
|
||||
|
||||
ret := C.metal_zk_poseidon2_merkle_layer(ctx, outFr, currentFr, C.uint32_t(n))
|
||||
if ret != C.METAL_ZK_SUCCESS {
|
||||
return nil, errors.New("Poseidon2 Merkle layer failed")
|
||||
}
|
||||
return output, nil
|
||||
}
|
||||
|
||||
// MerkleTree builds a complete Merkle tree from leaves using GPU.
|
||||
// Returns all internal nodes (num_leaves - 1 elements) with root at index 0.
|
||||
// Number of leaves must be a power of 2.
|
||||
func MerkleTree(leaves []Element) ([]Element, error) {
|
||||
n := len(leaves)
|
||||
if n == 0 || (n&(n-1)) != 0 {
|
||||
return nil, errors.New("leaves count must be a power of 2")
|
||||
}
|
||||
if err := initContext(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// Tree has n-1 internal nodes
|
||||
tree := make([]Element, n-1)
|
||||
leavesFr := (*C.Fr256)(unsafe.Pointer(&leaves[0]))
|
||||
treeFr := (*C.Fr256)(unsafe.Pointer(&tree[0]))
|
||||
|
||||
ret := C.metal_zk_poseidon2_merkle_tree(ctx, treeFr, leavesFr, C.uint32_t(n))
|
||||
if ret != C.METAL_ZK_SUCCESS {
|
||||
return nil, errors.New("Poseidon2 Merkle tree failed")
|
||||
}
|
||||
return tree, nil
|
||||
}
|
||||
|
||||
// MerkleRoot computes just the Merkle root from leaves.
|
||||
func MerkleRoot(leaves []Element) (Element, error) {
|
||||
var root Element
|
||||
tree, err := MerkleTree(leaves)
|
||||
if err != nil {
|
||||
return root, err
|
||||
}
|
||||
if len(tree) > 0 {
|
||||
root = tree[0]
|
||||
}
|
||||
return root, nil
|
||||
}
|
||||
|
||||
// VerifyProof verifies a single Merkle proof.
|
||||
// Returns true if the proof is valid.
|
||||
func VerifyProof(leaf, root Element, path []Element, indices []uint32) (bool, error) {
|
||||
results, err := BatchVerifyProofs([]Element{leaf}, []Element{root}, [][]Element{path}, [][]uint32{indices})
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
return results[0], nil
|
||||
}
|
||||
|
||||
// BatchVerifyProofs verifies multiple Merkle proofs in parallel on GPU.
|
||||
func BatchVerifyProofs(leaves, roots []Element, paths [][]Element, indices [][]uint32) ([]bool, error) {
|
||||
n := len(leaves)
|
||||
if n == 0 || n != len(roots) || n != len(paths) || n != len(indices) {
|
||||
return nil, errors.New("invalid input lengths")
|
||||
}
|
||||
if n > 0 && len(paths[0]) != len(indices[0]) {
|
||||
return nil, errors.New("path length mismatch")
|
||||
}
|
||||
if err := initContext(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
pathLen := len(paths[0])
|
||||
results := make([]C.int, n)
|
||||
|
||||
// Flatten paths and indices
|
||||
flatPaths := make([]Element, n*pathLen)
|
||||
flatIndices := make([]uint32, n*pathLen)
|
||||
for i := 0; i < n; i++ {
|
||||
copy(flatPaths[i*pathLen:], paths[i])
|
||||
copy(flatIndices[i*pathLen:], indices[i])
|
||||
}
|
||||
|
||||
leavesFr := (*C.Fr256)(unsafe.Pointer(&leaves[0]))
|
||||
rootsFr := (*C.Fr256)(unsafe.Pointer(&roots[0]))
|
||||
pathsFr := (*C.Fr256)(unsafe.Pointer(&flatPaths[0]))
|
||||
indicesPtr := (*C.uint32_t)(unsafe.Pointer(&flatIndices[0]))
|
||||
resultsPtr := (*C.int)(unsafe.Pointer(&results[0]))
|
||||
|
||||
ret := C.metal_zk_poseidon2_verify_proofs(
|
||||
ctx,
|
||||
resultsPtr,
|
||||
leavesFr,
|
||||
pathsFr,
|
||||
indicesPtr,
|
||||
rootsFr,
|
||||
C.uint32_t(n),
|
||||
C.uint32_t(pathLen),
|
||||
)
|
||||
if ret != C.METAL_ZK_SUCCESS {
|
||||
return nil, errors.New("Poseidon2 batch verify proofs failed")
|
||||
}
|
||||
|
||||
boolResults := make([]bool, n)
|
||||
for i := 0; i < n; i++ {
|
||||
boolResults[i] = results[i] != 0
|
||||
}
|
||||
return boolResults, nil
|
||||
}
|
||||
|
||||
// Commitment computes a Pedersen-style commitment using Poseidon2.
|
||||
// commitment = Poseidon2(value, blinding, salt)
|
||||
func Commitment(value, blinding, salt Element) (Element, error) {
|
||||
commitments, err := BatchCommitment([]Element{value}, []Element{blinding}, []Element{salt})
|
||||
if err != nil {
|
||||
return Element{}, err
|
||||
}
|
||||
return commitments[0], nil
|
||||
}
|
||||
|
||||
// BatchCommitment computes multiple commitments in parallel on GPU.
|
||||
func BatchCommitment(values, blindings, salts []Element) ([]Element, error) {
|
||||
n := len(values)
|
||||
if n == 0 || n != len(blindings) || n != len(salts) {
|
||||
return nil, errors.New("invalid input lengths")
|
||||
}
|
||||
if err := initContext(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
outputs := make([]Element, n)
|
||||
valuesFr := (*C.Fr256)(unsafe.Pointer(&values[0]))
|
||||
blindingsFr := (*C.Fr256)(unsafe.Pointer(&blindings[0]))
|
||||
saltsFr := (*C.Fr256)(unsafe.Pointer(&salts[0]))
|
||||
outFr := (*C.Fr256)(unsafe.Pointer(&outputs[0]))
|
||||
|
||||
ret := C.metal_zk_batch_commitment(ctx, outFr, valuesFr, blindingsFr, saltsFr, C.uint32_t(n))
|
||||
if ret != C.METAL_ZK_SUCCESS {
|
||||
return nil, errors.New("Poseidon2 batch commitment failed")
|
||||
}
|
||||
return outputs, nil
|
||||
}
|
||||
|
||||
// Nullifier computes a nullifier for spending a note.
|
||||
// nullifier = Poseidon2(key, commitment, index)
|
||||
func Nullifier(key, commitment, index Element) (Element, error) {
|
||||
nullifiers, err := BatchNullifier([]Element{key}, []Element{commitment}, []Element{index})
|
||||
if err != nil {
|
||||
return Element{}, err
|
||||
}
|
||||
return nullifiers[0], nil
|
||||
}
|
||||
|
||||
// BatchNullifier computes multiple nullifiers in parallel on GPU.
|
||||
func BatchNullifier(keys, commitments, indices []Element) ([]Element, error) {
|
||||
n := len(keys)
|
||||
if n == 0 || n != len(commitments) || n != len(indices) {
|
||||
return nil, errors.New("invalid input lengths")
|
||||
}
|
||||
if err := initContext(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
outputs := make([]Element, n)
|
||||
keysFr := (*C.Fr256)(unsafe.Pointer(&keys[0]))
|
||||
commitmentsFr := (*C.Fr256)(unsafe.Pointer(&commitments[0]))
|
||||
indicesFr := (*C.Fr256)(unsafe.Pointer(&indices[0]))
|
||||
outFr := (*C.Fr256)(unsafe.Pointer(&outputs[0]))
|
||||
|
||||
ret := C.metal_zk_batch_nullifier(ctx, outFr, keysFr, commitmentsFr, indicesFr, C.uint32_t(n))
|
||||
if ret != C.METAL_ZK_SUCCESS {
|
||||
return nil, errors.New("Poseidon2 batch nullifier failed")
|
||||
}
|
||||
return outputs, nil
|
||||
}
|
||||
|
||||
// Destroy releases the Metal ZK context.
|
||||
// Should be called when done with GPU operations.
|
||||
func Destroy() {
|
||||
if ctxReady && ctx != nil {
|
||||
C.metal_zk_destroy(ctx)
|
||||
ctx = nil
|
||||
ctxReady = false
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,56 @@
|
||||
// Copyright (C) 2025, Lux Industries, Inc. All rights reserved.
|
||||
// See the file LICENSE for licensing terms.
|
||||
|
||||
//go:build !luxgpu
|
||||
|
||||
// Package gpu provides GPU-accelerated Poseidon2 hash functions.
|
||||
// This is the stub implementation when GPU acceleration is not available.
|
||||
package gpu
|
||||
|
||||
import (
|
||||
"errors"
|
||||
)
|
||||
|
||||
// ErrNotAvailable is returned when GPU acceleration is not available.
|
||||
var ErrNotAvailable = errors.New("GPU acceleration not available")
|
||||
|
||||
// ElementSize is the size of a BN254 scalar field element in bytes (256-bit).
|
||||
const ElementSize = 32
|
||||
|
||||
// Element represents a 256-bit BN254 scalar field element.
|
||||
type Element [ElementSize]byte
|
||||
|
||||
// Available returns true if Metal GPU acceleration is available.
|
||||
func Available() bool {
|
||||
return false
|
||||
}
|
||||
|
||||
// Threshold returns the recommended batch size for GPU offload.
|
||||
func Threshold() uint32 {
|
||||
return 0
|
||||
}
|
||||
|
||||
// HashPair computes Poseidon2 hash of two elements (2-to-1 compression).
|
||||
func HashPair(left, right Element) (Element, error) {
|
||||
return Element{}, ErrNotAvailable
|
||||
}
|
||||
|
||||
// Hash computes Poseidon2 hash of variable-length input.
|
||||
func Hash(input []Element) (Element, error) {
|
||||
return Element{}, ErrNotAvailable
|
||||
}
|
||||
|
||||
// BatchHashPair computes multiple 2-to-1 compressions in parallel.
|
||||
func BatchHashPair(pairs [][2]Element) ([]Element, error) {
|
||||
return nil, ErrNotAvailable
|
||||
}
|
||||
|
||||
// MerkleRoot computes the Merkle tree root from leaves.
|
||||
func MerkleRoot(leaves []Element) (Element, error) {
|
||||
return Element{}, ErrNotAvailable
|
||||
}
|
||||
|
||||
// MerklePath computes the Merkle proof path for a leaf.
|
||||
func MerklePath(leaves []Element, index uint32) ([]Element, error) {
|
||||
return nil, ErrNotAvailable
|
||||
}
|
||||
@@ -0,0 +1,381 @@
|
||||
// Copyright (C) 2025, Lux Industries, Inc. All rights reserved.
|
||||
// See the file LICENSE for licensing terms.
|
||||
|
||||
//go:build cgo && darwin
|
||||
|
||||
// Package poseidon2 provides GPU-accelerated Poseidon2 hash functions via luxcpp/crypto.
|
||||
// Poseidon2 is a ZK-friendly hash function designed for BN254 scalar field.
|
||||
package poseidon2
|
||||
|
||||
/*
|
||||
#cgo CFLAGS: -I${SRCDIR}/../../../../luxcpp/crypto/include
|
||||
#cgo LDFLAGS: -L${SRCDIR}/../../../../luxcpp/crypto/build-local -lluxcrypto -framework Metal -framework Foundation
|
||||
|
||||
#include <stdint.h>
|
||||
#include <stdlib.h>
|
||||
#include <stdbool.h>
|
||||
#include "lux/crypto/metal_zk.h"
|
||||
|
||||
*/
|
||||
import "C"
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"unsafe"
|
||||
)
|
||||
|
||||
// ElementSize is the size of a BN254 scalar field element in bytes (256-bit).
|
||||
const ElementSize = 32
|
||||
|
||||
// Element represents a 256-bit BN254 scalar field element.
|
||||
// Stored as 4 x 64-bit limbs in little-endian order.
|
||||
type Element [ElementSize]byte
|
||||
|
||||
// context is the global Metal ZK context (initialized on first use).
|
||||
var (
|
||||
ctx *C.MetalZKContext
|
||||
ctxReady bool
|
||||
)
|
||||
|
||||
// initContext lazily initializes the Metal ZK context.
|
||||
func initContext() error {
|
||||
if ctxReady {
|
||||
return nil
|
||||
}
|
||||
ctx = C.metal_zk_init()
|
||||
if ctx == nil {
|
||||
return errors.New("Metal ZK initialization failed")
|
||||
}
|
||||
ctxReady = true
|
||||
return nil
|
||||
}
|
||||
|
||||
// GPUAvailable returns true if Metal GPU acceleration is available.
|
||||
func GPUAvailable() bool {
|
||||
return bool(C.metal_zk_available())
|
||||
}
|
||||
|
||||
// Threshold returns the recommended batch size for GPU offload.
|
||||
func Threshold() uint32 {
|
||||
return uint32(C.metal_zk_get_threshold(C.METAL_ZK_OP_POSEIDON2_HASH))
|
||||
}
|
||||
|
||||
// HashPair computes Poseidon2 hash of two elements (2-to-1 compression).
|
||||
// This is the core operation for Merkle tree construction.
|
||||
func HashPair(left, right Element) (Element, error) {
|
||||
var out Element
|
||||
if err := initContext(); err != nil {
|
||||
return out, err
|
||||
}
|
||||
|
||||
leftFr := (*C.Fr256)(unsafe.Pointer(&left[0]))
|
||||
rightFr := (*C.Fr256)(unsafe.Pointer(&right[0]))
|
||||
outFr := (*C.Fr256)(unsafe.Pointer(&out[0]))
|
||||
|
||||
ret := C.metal_zk_poseidon2_hash_pair(ctx, outFr, leftFr, rightFr, 1)
|
||||
if ret != C.METAL_ZK_SUCCESS {
|
||||
return out, errors.New("Poseidon2 hash pair failed")
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// BatchHashPair computes Poseidon2 hashes for multiple pairs on GPU.
|
||||
// More efficient than calling HashPair repeatedly for batch operations.
|
||||
func BatchHashPair(left, right []Element) ([]Element, error) {
|
||||
n := len(left)
|
||||
if n == 0 || n != len(right) {
|
||||
return nil, errors.New("invalid input lengths")
|
||||
}
|
||||
if err := initContext(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
outputs := make([]Element, n)
|
||||
|
||||
// Below threshold, process sequentially
|
||||
if uint32(n) < Threshold() || !GPUAvailable() {
|
||||
for i := 0; i < n; i++ {
|
||||
out, err := HashPair(left[i], right[i])
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
outputs[i] = out
|
||||
}
|
||||
return outputs, nil
|
||||
}
|
||||
|
||||
// GPU batch processing
|
||||
leftFr := (*C.Fr256)(unsafe.Pointer(&left[0]))
|
||||
rightFr := (*C.Fr256)(unsafe.Pointer(&right[0]))
|
||||
outFr := (*C.Fr256)(unsafe.Pointer(&outputs[0]))
|
||||
|
||||
ret := C.metal_zk_poseidon2_hash_pair(ctx, outFr, leftFr, rightFr, C.uint32_t(n))
|
||||
if ret != C.METAL_ZK_SUCCESS {
|
||||
return nil, errors.New("Poseidon2 batch hash pair failed")
|
||||
}
|
||||
return outputs, nil
|
||||
}
|
||||
|
||||
// MerkleLayer computes one layer of a Merkle tree from leaf pairs.
|
||||
// Input size must be even. Output size is input size / 2.
|
||||
func MerkleLayer(current []Element) ([]Element, error) {
|
||||
n := len(current)
|
||||
if n == 0 || n%2 != 0 {
|
||||
return nil, errors.New("layer size must be even")
|
||||
}
|
||||
if err := initContext(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
output := make([]Element, n/2)
|
||||
currentFr := (*C.Fr256)(unsafe.Pointer(¤t[0]))
|
||||
outFr := (*C.Fr256)(unsafe.Pointer(&output[0]))
|
||||
|
||||
ret := C.metal_zk_poseidon2_merkle_layer(ctx, outFr, currentFr, C.uint32_t(n))
|
||||
if ret != C.METAL_ZK_SUCCESS {
|
||||
return nil, errors.New("Poseidon2 Merkle layer failed")
|
||||
}
|
||||
return output, nil
|
||||
}
|
||||
|
||||
// MerkleTree builds a complete Merkle tree from leaves using GPU.
|
||||
// Returns all internal nodes (num_leaves - 1 elements) with root at index 0.
|
||||
// Number of leaves must be a power of 2.
|
||||
func MerkleTree(leaves []Element) ([]Element, error) {
|
||||
n := len(leaves)
|
||||
if n == 0 || (n&(n-1)) != 0 {
|
||||
return nil, errors.New("leaves count must be a power of 2")
|
||||
}
|
||||
if err := initContext(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// Tree has n-1 internal nodes
|
||||
tree := make([]Element, n-1)
|
||||
leavesFr := (*C.Fr256)(unsafe.Pointer(&leaves[0]))
|
||||
treeFr := (*C.Fr256)(unsafe.Pointer(&tree[0]))
|
||||
|
||||
ret := C.metal_zk_poseidon2_merkle_tree(ctx, treeFr, leavesFr, C.uint32_t(n))
|
||||
if ret != C.METAL_ZK_SUCCESS {
|
||||
return nil, errors.New("Poseidon2 Merkle tree failed")
|
||||
}
|
||||
return tree, nil
|
||||
}
|
||||
|
||||
// MerkleRoot computes just the Merkle root from leaves.
|
||||
func MerkleRoot(leaves []Element) (Element, error) {
|
||||
var root Element
|
||||
tree, err := MerkleTree(leaves)
|
||||
if err != nil {
|
||||
return root, err
|
||||
}
|
||||
if len(tree) > 0 {
|
||||
root = tree[0]
|
||||
}
|
||||
return root, nil
|
||||
}
|
||||
|
||||
// VerifyProof verifies a single Merkle proof.
|
||||
// Returns true if the proof is valid.
|
||||
func VerifyProof(leaf, root Element, path []Element, indices []uint32) (bool, error) {
|
||||
results, err := BatchVerifyProofs([]Element{leaf}, []Element{root}, [][]Element{path}, [][]uint32{indices})
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
return results[0], nil
|
||||
}
|
||||
|
||||
// BatchVerifyProofs verifies multiple Merkle proofs in parallel on GPU.
|
||||
func BatchVerifyProofs(leaves, roots []Element, paths [][]Element, indices [][]uint32) ([]bool, error) {
|
||||
n := len(leaves)
|
||||
if n == 0 || n != len(roots) || n != len(paths) || n != len(indices) {
|
||||
return nil, errors.New("invalid input lengths")
|
||||
}
|
||||
if n > 0 && len(paths[0]) != len(indices[0]) {
|
||||
return nil, errors.New("path length mismatch")
|
||||
}
|
||||
if err := initContext(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
pathLen := len(paths[0])
|
||||
results := make([]C.int, n)
|
||||
|
||||
// Flatten paths and indices
|
||||
flatPaths := make([]Element, n*pathLen)
|
||||
flatIndices := make([]uint32, n*pathLen)
|
||||
for i := 0; i < n; i++ {
|
||||
copy(flatPaths[i*pathLen:], paths[i])
|
||||
copy(flatIndices[i*pathLen:], indices[i])
|
||||
}
|
||||
|
||||
leavesFr := (*C.Fr256)(unsafe.Pointer(&leaves[0]))
|
||||
rootsFr := (*C.Fr256)(unsafe.Pointer(&roots[0]))
|
||||
pathsFr := (*C.Fr256)(unsafe.Pointer(&flatPaths[0]))
|
||||
indicesPtr := (*C.uint32_t)(unsafe.Pointer(&flatIndices[0]))
|
||||
resultsPtr := (*C.int)(unsafe.Pointer(&results[0]))
|
||||
|
||||
ret := C.metal_zk_poseidon2_verify_proofs(
|
||||
ctx,
|
||||
resultsPtr,
|
||||
leavesFr,
|
||||
pathsFr,
|
||||
indicesPtr,
|
||||
rootsFr,
|
||||
C.uint32_t(n),
|
||||
C.uint32_t(pathLen),
|
||||
)
|
||||
if ret != C.METAL_ZK_SUCCESS {
|
||||
return nil, errors.New("Poseidon2 batch verify proofs failed")
|
||||
}
|
||||
|
||||
boolResults := make([]bool, n)
|
||||
for i := 0; i < n; i++ {
|
||||
boolResults[i] = results[i] != 0
|
||||
}
|
||||
return boolResults, nil
|
||||
}
|
||||
|
||||
// Commitment computes a Pedersen-style commitment using Poseidon2.
|
||||
// commitment = Poseidon2(value, blinding, salt)
|
||||
func Commitment(value, blinding, salt Element) (Element, error) {
|
||||
commitments, err := BatchCommitment([]Element{value}, []Element{blinding}, []Element{salt})
|
||||
if err != nil {
|
||||
return Element{}, err
|
||||
}
|
||||
return commitments[0], nil
|
||||
}
|
||||
|
||||
// BatchCommitment computes multiple commitments in parallel on GPU.
|
||||
func BatchCommitment(values, blindings, salts []Element) ([]Element, error) {
|
||||
n := len(values)
|
||||
if n == 0 || n != len(blindings) || n != len(salts) {
|
||||
return nil, errors.New("invalid input lengths")
|
||||
}
|
||||
if err := initContext(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
outputs := make([]Element, n)
|
||||
valuesFr := (*C.Fr256)(unsafe.Pointer(&values[0]))
|
||||
blindingsFr := (*C.Fr256)(unsafe.Pointer(&blindings[0]))
|
||||
saltsFr := (*C.Fr256)(unsafe.Pointer(&salts[0]))
|
||||
outFr := (*C.Fr256)(unsafe.Pointer(&outputs[0]))
|
||||
|
||||
ret := C.metal_zk_batch_commitment(ctx, outFr, valuesFr, blindingsFr, saltsFr, C.uint32_t(n))
|
||||
if ret != C.METAL_ZK_SUCCESS {
|
||||
return nil, errors.New("Poseidon2 batch commitment failed")
|
||||
}
|
||||
return outputs, nil
|
||||
}
|
||||
|
||||
// Nullifier computes a nullifier for spending a note.
|
||||
// nullifier = Poseidon2(key, commitment, index)
|
||||
func Nullifier(key, commitment, index Element) (Element, error) {
|
||||
nullifiers, err := BatchNullifier([]Element{key}, []Element{commitment}, []Element{index})
|
||||
if err != nil {
|
||||
return Element{}, err
|
||||
}
|
||||
return nullifiers[0], nil
|
||||
}
|
||||
|
||||
// BatchNullifier computes multiple nullifiers in parallel on GPU.
|
||||
func BatchNullifier(keys, commitments, indices []Element) ([]Element, error) {
|
||||
n := len(keys)
|
||||
if n == 0 || n != len(commitments) || n != len(indices) {
|
||||
return nil, errors.New("invalid input lengths")
|
||||
}
|
||||
if err := initContext(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
outputs := make([]Element, n)
|
||||
keysFr := (*C.Fr256)(unsafe.Pointer(&keys[0]))
|
||||
commitmentsFr := (*C.Fr256)(unsafe.Pointer(&commitments[0]))
|
||||
indicesFr := (*C.Fr256)(unsafe.Pointer(&indices[0]))
|
||||
outFr := (*C.Fr256)(unsafe.Pointer(&outputs[0]))
|
||||
|
||||
ret := C.metal_zk_batch_nullifier(ctx, outFr, keysFr, commitmentsFr, indicesFr, C.uint32_t(n))
|
||||
if ret != C.METAL_ZK_SUCCESS {
|
||||
return nil, errors.New("Poseidon2 batch nullifier failed")
|
||||
}
|
||||
return outputs, nil
|
||||
}
|
||||
|
||||
// Destroy releases the Metal ZK context.
|
||||
// Should be called when done with GPU operations.
|
||||
func Destroy() {
|
||||
if ctxReady && ctx != nil {
|
||||
C.metal_zk_destroy(ctx)
|
||||
ctx = nil
|
||||
ctxReady = false
|
||||
}
|
||||
}
|
||||
|
||||
// Hash computes Poseidon2 hash of multiple elements (1-16) using Merkle-Damgard construction.
|
||||
// This mirrors the gnark-crypto Poseidon2 Merkle-Damgard hasher behavior.
|
||||
// For 2 elements, it uses HashPair directly. For more, it chains hashes.
|
||||
func Hash(elements []Element) (Element, error) {
|
||||
n := len(elements)
|
||||
if n == 0 {
|
||||
return Element{}, errors.New("empty input")
|
||||
}
|
||||
if n > 16 {
|
||||
return Element{}, errors.New("too many elements: maximum 16")
|
||||
}
|
||||
|
||||
// Single element: hash with zero element
|
||||
if n == 1 {
|
||||
return HashPair(elements[0], Element{})
|
||||
}
|
||||
|
||||
// Two elements: direct hash pair
|
||||
if n == 2 {
|
||||
return HashPair(elements[0], elements[1])
|
||||
}
|
||||
|
||||
// More elements: Merkle-Damgard construction
|
||||
// H(e0, e1) -> state
|
||||
// H(state, e2) -> state
|
||||
// ... continue until all elements consumed
|
||||
state, err := HashPair(elements[0], elements[1])
|
||||
if err != nil {
|
||||
return Element{}, err
|
||||
}
|
||||
|
||||
for i := 2; i < n; i++ {
|
||||
state, err = HashPair(state, elements[i])
|
||||
if err != nil {
|
||||
return Element{}, err
|
||||
}
|
||||
}
|
||||
|
||||
return state, nil
|
||||
}
|
||||
|
||||
// HashFromBytes computes Poseidon2 hash from raw bytes.
|
||||
// Input must be a multiple of 32 bytes (1-16 field elements).
|
||||
// This is the main entry point for precompile use.
|
||||
func HashFromBytes(input []byte) ([]byte, error) {
|
||||
if len(input) == 0 || len(input)%ElementSize != 0 {
|
||||
return nil, errors.New("invalid input length: must be multiple of 32 bytes")
|
||||
}
|
||||
|
||||
numElements := len(input) / ElementSize
|
||||
if numElements > 16 {
|
||||
return nil, errors.New("too many inputs: maximum 16 field elements")
|
||||
}
|
||||
|
||||
// Parse into Elements
|
||||
elements := make([]Element, numElements)
|
||||
for i := 0; i < numElements; i++ {
|
||||
copy(elements[i][:], input[i*ElementSize:(i+1)*ElementSize])
|
||||
}
|
||||
|
||||
result, err := Hash(elements)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return result[:], nil
|
||||
}
|
||||
Reference in New Issue
Block a user