From d11ee38d5b780da24cd62f5d355ae18be81b1daf Mon Sep 17 00:00:00 2001
From: Giuseppe <8973725+Eagle941@users.noreply.github.com>
Date: Wed, 29 Nov 2023 18:17:45 +0000
Subject: [PATCH] Added gnark version switch

---
 extractor/extractor.go   |  8 ++++++++
 extractor/interface.go   |  4 ++--
 extractor/lean_export.go | 25 ++++++++++++++++++++-----
 3 files changed, 30 insertions(+), 7 deletions(-)

diff --git a/extractor/extractor.go b/extractor/extractor.go
index 3fa39dc..0cba6b8 100644
--- a/extractor/extractor.go
+++ b/extractor/extractor.go
@@ -165,7 +165,15 @@ type ExArg struct {
 	Type ExArgType
 }
 
+type GnarkVersion int
+
+const (
+	Gnark8 GnarkVersion = iota
+	Gnark9
+)
+
 type ExCircuit struct {
+	Version GnarkVersion
 	Inputs  []ExArg
 	Gadgets []ExGadget
 	Code    []App
diff --git a/extractor/interface.go b/extractor/interface.go
index 82fb80c..c41b88d 100644
--- a/extractor/interface.go
+++ b/extractor/interface.go
@@ -42,8 +42,8 @@ func GenerateLeanCircuit(name string, ext ExtractorApi, circuit ExtractorCircuit
 	return fmt.Sprintf("def %s %s: Prop :=\n%s", name, genArgs(extractorCircuit.Inputs), genCircuitBody(extractorCircuit))
 }
 
-func GenerateLeanCircuits(namespace string, ext ExtractorApi, circuits []string) string {
-	prelude := exportPrelude(namespace, ext.GetField().ScalarField())
+func GenerateLeanCircuits(namespace string, ext ExtractorApi, circuits []string, gateVersion GnarkVersion) string {
+	prelude := exportPrelude(namespace, ext.GetField().ScalarField(), genGateVersion(gateVersion))
 	gadgets := exportGadgets(ext.GetGadgets())
 	footer := exportFooter(namespace)
 	return fmt.Sprintf("%s\n\n%s\n\n%s\n\n%s", prelude, gadgets, strings.Join(circuits, "\n\n"), footer)
diff --git a/extractor/lean_export.go b/extractor/lean_export.go
index 5592499..fff8bfe 100644
--- a/extractor/lean_export.go
+++ b/extractor/lean_export.go
@@ -17,10 +17,24 @@ func isWhitespacePresent(input string) bool {
 	return regexp.MustCompile(`\s`).MatchString(input)
 }
 
+// genGateVersion returns the string to be used to make Lean use the correct
+// version of the Gates (depending on Gnark version in use)
+func genGateVersion(v GnarkVersion) string {
+	name := "unknown"
+	switch v {
+	case Gnark8:
+		name = "GatesGnark_8"
+	case Gnark9:
+		name = "GatesGnark_9"
+	}
+
+	return name
+}
+
 // exportPrelude generates the string to put at the beginning of the
 // autogenerated Lean4 code. It includes the relevant `provenZK` library
 // import.
-func exportPrelude(name string, order *big.Int) string {
+func exportPrelude(name string, order *big.Int, gateVersion string) string {
 	trimmedName := strings.TrimSpace(name)
 	if isWhitespacePresent(trimmedName) {
 		panic("Whitespace isn't allowed in namespace tag")
@@ -34,7 +48,8 @@ namespace %s
 
 def Order : ℕ := 0x%s
 variable [Fact (Nat.Prime Order)]
-abbrev F := ZMod Order`, trimmedName, order.Text(16))
+abbrev F := ZMod Order
+abbrev Gates := %s Order`, trimmedName, order.Text(16), gateVersion)
 
 	return s
 }
@@ -92,15 +107,15 @@ func exportCircuitFunction(name string, circuit ExCircuit) string {
 
 // exportCircuit generates the `circuit` function in Lean
 func exportCircuit(circuit ExCircuit, name string) string {
-	prelude := exportPrelude(name, circuit.Field.ScalarField())
+	prelude := exportPrelude(name, circuit.Field.ScalarField(), genGateVersion(circuit.Version))
 	gadgets := exportGadgets(circuit.Gadgets)
 	circ := exportCircuitFunction("circuit", circuit)
 	footer := exportFooter(name)
 	return fmt.Sprintf("%s\n\n%s\n\n%s\n\n%s", prelude, gadgets, circ, footer)
 }
 
-func ExportGadgetsOnly(namespace string, exGadgets []ExGadget, field ecc.ID) string {
-	prelude := exportPrelude(namespace, field.ScalarField())
+func ExportGadgetsOnly(namespace string, exGadgets []ExGadget, field ecc.ID, gateVersion GnarkVersion) string {
+	prelude := exportPrelude(namespace, field.ScalarField(), genGateVersion(gateVersion))
 	gadgets := exportGadgets(exGadgets)
 	footer := exportFooter(namespace)
 	return fmt.Sprintf("%s\n\n%s\n\n%s", prelude, gadgets, footer)