Zksync separate function

contracts/.gitignore

@@ -2,6 +2,7 @@ artifacts
 cache
 node_modules
 solc
+zksolc
 abi
 coverage
 coverage.json

contracts/GNUmakefile

@@ -37,6 +37,16 @@ pnpmdep: ## Install solidity contract dependencies through pnpm
 abigen: ## Build & install abigen.
 	../tools/bin/build_abigen
 
+
+# env var example
+#		export ZKSOLC_VERSION=1.5.6
+#		make setup_zksolc
+# make call example
+# 		make ZKSOLC_VERSION=1.5.6 setup_zksolc
+.PHONY: setup_zksolc
+setup_zksolc: ## Download and install zksolc
+	../tools/bin/setup_zksolc.sh $(ZKSOLC_VERSION)
+
 .PHONY: mockery
 mockery: $(mockery) ## Install mockery.
 	go install github.com/vektra/mockery/v2@v2.43.2
@@ -76,15 +86,15 @@ ccip-lcov:
 # make call example
 # 		make FOUNDRY_PROFILE=llo-feeds wrappers
 .PHONY: wrappers
-wrappers: pnpmdep mockery abigen ## Recompiles solidity contracts and their go wrappers.
+wrappers: pnpmdep mockery abigen setup_zksolc ## Recompiles solidity contracts and their go wrappers.
 	./scripts/native_solc_compile_all_$(FOUNDRY_PROFILE)
 	go generate ../core/gethwrappers/$(FOUNDRY_PROFILE)
 
 # This call generates all gethwrappers for all products. It does so based on the
 # assumption that native_solc_compile_all contains sub-calls to each product, and
 # go_generate does the same.
 .PHONY: wrappers-all
-wrappers-all: pnpmdep mockery abigen ## Recompiles solidity contracts and their go wrappers.
+wrappers-all: pnpmdep mockery abigen setup_zksolc ## Recompiles solidity contracts and their go wrappers.
 	# go_generate contains a call to  compile all contracts before generating wrappers
 	go generate ../core/gethwrappers/go_generate.go
 
@@ -108,4 +118,4 @@ help:
 	@echo "       \/     \/     \/        \/             \/     \/"
 	@echo ""
 	@grep -E '^[a-zA-Z_-]+:.*?## .*$$' $(MAKEFILE_LIST) | sort | \
-	awk 'BEGIN {FS = ":.*?## "}; {printf "\033[36m%-30s\033[0m %s\n", $$1, $$2}'
+	awk 'BEGIN {FS = ":.*?## "}; {printf "\033[36m%-30s\033[0m %s\n", $$1, $$2}'

contracts/scripts/native_solc_compile_all_ccip

@@ -9,6 +9,7 @@ echo " └───────────────────────
 SOLC_VERSION="0.8.24"
 OPTIMIZE_RUNS=26000
 OPTIMIZE_RUNS_OFFRAMP=500
+ZKSOLC_VERSION="1.5.3"
 
 
 SCRIPTPATH="$( cd "$(dirname "$0")" >/dev/null 2>&1 ; pwd -P )"
@@ -38,6 +39,13 @@ compileContract () {
       --bin-runtime --hashes --metadata --metadata-literal --combined-json abi,hashes,metadata,srcmap,srcmap-runtime \
       --evm-version paris \
       "$ROOT"/contracts/src/v0.8/"$1"
+
+  zksolc --overwrite -O3 --metadata-hash none \
+      -o "$ROOT"/contracts/zksolc/v$ZKSOLC_VERSION/"$contract" \
+      --bin --allow-paths "$ROOT"/contracts/src/v0.8 \
+      --metadata-literal \
+      --evm-version paris \
+      "$ROOT"/contracts/src/v0.8/"$1"  
 }
 
 
@@ -78,13 +86,13 @@ compileContract ccip/test/helpers/receivers/MaybeRevertMessageReceiver.sol
 compileContract ccip/test/helpers/MultiOCR3Helper.sol
 compileContract ccip/test/mocks/MockE2EUSDCTokenMessenger.sol
 compileContract ccip/test/mocks/MockE2EUSDCTransmitter.sol
-compileContract ccip/test/WETH9.sol
+# compileContract ccip/test/WETH9.sol (error)
 compileContract ccip/test/helpers/CCIPReaderTester.sol
 
 # Encoding Utils
-compileContract ccip/interfaces/encodingutils/ICCIPEncodingUtils.sol
+# compileContract ccip/interfaces/encodingutils/ICCIPEncodingUtils.sol (error)
 
 # Customer contracts
 compileContract ccip/pools/USDC/USDCTokenPool.sol
 
-compileContract tests/MockV3Aggregator.sol
+compileContract tests/MockV3Aggregator.sol

core/gethwrappers/abigen.go

@@ -6,7 +6,9 @@
 	"go/ast"
 	"go/format"
 	"go/parser"
+	"go/printer"
 	"go/token"
+	"log"
 	"os"
 	"os/exec"
 	"path/filepath"
@@ -28,7 +30,7 @@
 // AbigenArgs is the arguments to the abigen executable. E.g., Bin is the -bin
 // arg.
 type AbigenArgs struct {
-	Bin, ABI, Out, Type, Pkg string
+	Bin, ABI, Out, Type, Pkg, ZkBinPath string
 }
 
 // Abigen calls Abigen  with the given arguments
@@ -72,6 +74,9 @@
 	}
 
 	ImproveAbigenOutput(a.Out, a.ABI)
+	if a.ZkBinPath != "" {
+		ImproveAbigenOutputZKs(a.Out, a.ZkBinPath)
+	}
 }
 
 func ImproveAbigenOutput(path string, abiPath string) {
@@ -466,3 +471,269 @@
 func addHeader(code []byte) []byte {
 	return utils.ConcatBytes([]byte(headerComment), code)
 }
+
+// ZK stack logic
+func ImproveAbigenOutputZKs(path string, zkBinPath string) {
+
+	bs, err := os.ReadFile(path)
+	if err != nil {
+		Exit("Error while improving abigen output", err)
+	}
+
+	fset, fileNode := parseFile(bs)
+
+	contractName := getContractName(fileNode)
+
+	zkByteCode, err := os.ReadFile(zkBinPath)
+	if err != nil {
+		Exit("Error while improving abigen output", err)
+	}
+	zkHexString := string(zkByteCode)
+
+	// add zksync binary to the wrapper
+	fileNode = addZKSyncBin(fileNode, contractName, zkHexString)
+
+	// add zksync logic to the deploy function
+	fileNode = updateDeployMethod(contractName, fset, fileNode)
+
+	bs = generateCode(fset, fileNode)
+
+	err = os.WriteFile(path, bs, 0600)
+	if err != nil {
+		Exit("Error while writing improved abigen source", err)
+	}
+}
+
+// add zksync binary to the wrapper
+func addZKSyncBin(fileNode *ast.File, contractName string, zkHexString string) *ast.File {
+	// zksync
+	newVarSpec := &ast.ValueSpec{
+		Names: []*ast.Ident{ast.NewIdent(fmt.Sprintf("%sZKBin", contractName))},
+		// Type:  ast.NewIdent("string"),
+		Values: []ast.Expr{
+			&ast.BasicLit{
+				// Kind:  token.STRING,
+				Value: fmt.Sprintf("(\"0x%s\")", zkHexString),
+			},
+		},
+	}
+	newVarDecl := &ast.GenDecl{
+		Tok:   token.VAR,
+		Specs: []ast.Spec{newVarSpec},
+	}
+
+	// Insert the new variable declaration at the top of the file (before existing functions)
+	fileNode.Decls = append(fileNode.Decls, newVarDecl)
+	return fileNode
+}
+
+// add zksync logic to the deploy function
+func updateDeployMethod(contractName string, fset *token.FileSet, fileNode *ast.File) *ast.File {
+
+	return astutil.Apply(fileNode, func(cursor *astutil.Cursor) bool {
+		x, is := cursor.Node().(*ast.FuncDecl)
+		if !is {
+			return true
+		} else if x.Name.Name != "Deploy"+contractName {
+			return false
+		}
+
+		// only add this import if Deploy method found
+		astutil.AddImport(fset, fileNode, "github.com/smartcontractkit/chainlink/v2/core/gethwrappers/generated")
+
+		fmt.Println("here")
+
+		// Extract the parameters from the existing function x
+		paramList := getConstructorParams(x.Type.Params.List)
+		paramListWithTypes := getConstructorParamsWithTypes(x.Type.Params.List)
+		fmt.Println("paramListWithTypes: ", paramListWithTypes)
+		zkSyncBlock := getZKSyncBlock(contractName, paramList, paramListWithTypes)
+		fmt.Println("zkSyncBlock: ", zkSyncBlock)
+		node, err := parser.ParseFile(token.NewFileSet(), "", zkSyncBlock, parser.SkipObjectResolution)
+		if err != nil {
+			return false
+		}
+
+		var newFuncDecl *ast.FuncDecl
+		for _, decl := range node.Decls {
+			if funcDecl, ok := decl.(*ast.FuncDecl); ok {
+				newFuncDecl = funcDecl
+				break
+			}
+		}
+
+		// Add the standalone function to the file
+		fileNode.Decls = append(fileNode.Decls, newFuncDecl)
+
+		return false
+	}, nil).(*ast.File)
+}
+
+func parseFuncDecl(funcCode string) (*ast.FuncDecl, error) {
+	// Parse the function string into a Go AST
+	node, err := parser.ParseFile(token.NewFileSet(), "", funcCode, parser.SkipObjectResolution)
+	if err != nil {
+		return nil, fmt.Errorf("error parsing function: %w", err)
+	}
+
+	// Extract the first declaration from the file node
+	for _, decl := range node.Decls {
+		if funcDecl, ok := decl.(*ast.FuncDecl); ok {
+			return funcDecl, nil
+		}
+	}
+
+	return nil, fmt.Errorf("no function declaration found")
+}
+
+var zkcode = `
+	package main
+	func Deploy%sZK(auth *bind.TransactOpts, backend bind.ContractBackend, %constructorParams) (common.Address, *generated.Transaction, *%s, error) {
+		parsed, err := %sMetaData.GetAbi()
+		if err != nil {
+			return common.Address{}, nil, nil, err
+		}
+		if parsed == nil {
+			return common.Address{}, nil, nil, errors.New("GetABI returned nil")
+		}
+		address, ethTx, contract, err := generated.DeployContract(auth, parsed, common.FromHex(%sZKBin), backend, %deployParams)
+		if err != nil {
+			return common.Address{}, nil, nil, err
+		}
+		return address, ethTx, &%s{address: address, abi: *parsed, %sCaller: %sCaller{contract: contract}, %sTransactor: %sTransactor{contract: contract}, %sFilterer: %sFilterer{contract: contract}}, nil
+	}
+`
+
+// get the `if zksync()` block
+func getZKSyncBlock(contractName, paramList, paramListWithTypes string) string {
+	// zkSyncBlock := `if generated_zks.IsZKSync(backend) {
+	// 			address, ethTx, contractBind, _ := generated_zks.DeployContract(auth, parsed, common.FromHex(%sZKBin), backend, %params)
+	// 			contractReturn := &%s{address: address, abi: *parsed, %sCaller: %sCaller{contract: contractBind}, %sTransactor: %sTransactor{contract: contractBind},%sFilterer: %sFilterer{contract: contractBind}}
+	// 			return address, ethTx, contractReturn, err
+	// 	}`
+	zkSyncBlock := zkcode
+	zkSyncBlock = strings.ReplaceAll(zkSyncBlock, "%s", contractName)
+	zkSyncBlock = strings.ReplaceAll(zkSyncBlock, "%constructorParams", paramListWithTypes)
+	zkSyncBlock = strings.ReplaceAll(zkSyncBlock, "%deployParams", paramList)
+	return strings.ReplaceAll(zkSyncBlock, "%s", contractName)
+}
+
+// Extract the parameters for constructor function
+func getConstructorParams(contstructorParams []*ast.Field) string {
+	params := []string{}
+	for i, param := range contstructorParams {
+		if i > 1 { // Skip auth and backend
+			for _, name := range param.Names {
+				params = append(params, name.Name)
+			}
+		}
+	}
+	paramList := strings.Join(params, ", ")
+	return paramList
+}
+
+func getConstructorParamsWithTypes(constructorParams []*ast.Field) string {
+	params := []string{}
+	for i, param := range constructorParams {
+		if i > 1 { // Skip auth and backend
+			for _, name := range param.Names {
+				paramWithType := fmt.Sprintf("%s %s", name.Name, exprToString(param.Type))
+				params = append(params, paramWithType)
+			}
+		}
+	}
+	return strings.Join(params, ", ")
+}
+
+func exprToString(expr ast.Expr) string {
+	var buf bytes.Buffer
+	err := printer.Fprint(&buf, token.NewFileSet(), expr)
+	if err != nil {
+		log.Fatalf("Failed to convert AST expression to string: %v", err)
+	}
+	return buf.String()
+}
+
+// insert the `if zksync()` block
+func addZKSyncBlock(x ast.FuncDecl, zkSyncBlock string) ast.FuncDecl {
+	for i, stmt := range x.Body.List {
+
+		ifStmt, ok := stmt.(*ast.IfStmt)
+		if !ok {
+			continue
+		}
+		binaryExpr, ok := ifStmt.Cond.(*ast.BinaryExpr)
+		if !ok {
+			continue
+		}
+		if ident, ok := binaryExpr.X.(*ast.Ident); ok && ident.Name == "parsed" {
+			// Creating new statement to insert
+			newStmt := &ast.ExprStmt{
+				X: &ast.BasicLit{
+					Kind:  token.STRING,
+					Value: zkSyncBlock,
+				},
+			}
+
+			// Insert the new statement after the current statement
+			x.Body.List = append(x.Body.List[:i+1], append([]ast.Stmt{newStmt}, x.Body.List[i+1:]...)...)
+			break
+		}
+	}
+	return x
+}
+
+// convert *types.Transaction to *generated_zks.Transaction
+func updateTxReturnType(x ast.FuncDecl) {
+	x.Type.Results.List[1].Type = &ast.StarExpr{
+		X: &ast.SelectorExpr{
+			X:   &ast.Ident{Name: "generated_zks"},
+			Sel: &ast.Ident{Name: "Transaction"},
+		},
+	}
+}
+
+// convert tx to &Transaction{Transaction: tx, Hash_zks: tx.Hash()}
+func updateReturnStmt(x ast.FuncDecl) {
+	for _, stmt := range x.Body.List {
+		returnStmt, is := stmt.(*ast.ReturnStmt)
+		if !is {
+			continue
+		}
+		if len(returnStmt.Results) < 3 {
+			continue
+		}
+
+		txExpr, ok := returnStmt.Results[1].(*ast.Ident)
+		if !ok {
+			return
+		}
+		if txExpr.Name != "tx" {
+			return
+		}
+
+		txField := &ast.KeyValueExpr{
+			Key:   ast.NewIdent("Transaction"),
+			Value: ast.NewIdent("tx"),
+		}
+
+		hashField := &ast.KeyValueExpr{
+			Key: ast.NewIdent("Hash_zks"),
+			Value: &ast.CallExpr{
+				Fun: &ast.SelectorExpr{
+					X:   ast.NewIdent("tx"),
+					Sel: ast.NewIdent("Hash"),
+				},
+			},
+		}
+		newRet := &ast.CompositeLit{
+			Type: &ast.SelectorExpr{
+				X:   ast.NewIdent("generated_zks"),
+				Sel: ast.NewIdent("Transaction"),
+			},
+			Elts: []ast.Expr{txField, hashField},
+		}
+		pointerRet := &ast.UnaryExpr{Op: token.AND, X: newRet}
+		returnStmt.Results[1] = pointerRet
+	}
+}