CCIP gethwrappers with zkstack support

contracts/.gitignore

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

contracts/GNUmakefile

@@ -98,6 +98,12 @@ artifact-generate:
 	./scripts/generate_compiled_json_ccip.sh $(contract) $(solcversion) $(artifactpath)
 
 
+.PHONY: wrappers-ccip-zksupport
+wrappers-ccip-zksupport: pnpmdep mockery abigen ## Recompiles solidity contracts for ccip and their go wrappers with zksync support
+	./scripts/native_solc_compile_all_ccip_zks
+	go generate ../core/gethwrappers/ccip/go_generate_zks.go
+
+
 help:
 	@echo ""
 	@echo "         .__           .__       .__  .__        __"

contracts/scripts/native_solc_compile_all_ccip_zks

@@ -0,0 +1,98 @@
+#!/usr/bin/env bash
+
+set -e
+
+echo " ┌──────────────────────────────────────────────┐"
+echo " │          Compiling CCIP contracts...         │"
+echo " └──────────────────────────────────────────────┘"
+
+SOLC_VERSION="0.8.24"
+OPTIMIZE_RUNS=26000
+OPTIMIZE_RUNS_OFFRAMP=500
+ZKSOLC_VERSION="1.5.6"
+
+
+SCRIPTPATH="$( cd "$(dirname "$0")" >/dev/null 2>&1 ; pwd -P )"
+python3 -m pip install --require-hashes -r "$SCRIPTPATH"/requirements.txt
+solc-select install $SOLC_VERSION
+solc-select use $SOLC_VERSION
+export SOLC_VERSION=$SOLC_VERSION
+
+ROOT="$( cd "$(dirname "$0")" >/dev/null 2>&1 ; cd ../../ && pwd -P )"
+
+compileContract () {
+  local contract
+  contract=$(basename "$1" ".sol")
+
+  local optimize_runs=$OPTIMIZE_RUNS
+
+  case $1 in
+    "ccip/offRamp/OffRamp.sol")
+    echo "OffRamp uses $OPTIMIZE_RUNS_OFFRAMP optimizer runs."
+      optimize_runs=$OPTIMIZE_RUNS_OFFRAMP
+      ;;
+  esac
+
+  solc --overwrite --optimize --optimize-runs $optimize_runs --metadata-hash none \
+      -o "$ROOT"/contracts/solc/v$SOLC_VERSION/"$contract" \
+      --abi --bin --allow-paths "$ROOT"/contracts/src/v0.8 \
+      --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 --metadata-literal \
+      --evm-version paris \
+      "$ROOT"/contracts/src/v0.8/"$1"  
+}
+
+
+# Solc produces and overwrites intermediary contracts.
+# Contracts should be ordered in reverse-import-complexity-order to minimize overwrite risks.
+compileContract ccip/offRamp/OffRamp.sol
+compileContract ccip/FeeQuoter.sol
+compileContract ccip/onRamp/OnRamp.sol
+compileContract ccip/applications/PingPongDemo.sol
+compileContract ccip/applications/EtherSenderReceiver.sol
+compileContract ccip/rmn/RMNRemote.sol
+compileContract ccip/rmn/RMNHome.sol
+compileContract ccip/rmn/ARMProxy.sol
+compileContract ccip/MultiAggregateRateLimiter.sol
+compileContract ccip/Router.sol
+compileContract ccip/tokenAdminRegistry/TokenAdminRegistry.sol
+compileContract ccip/tokenAdminRegistry/RegistryModuleOwnerCustom.sol
+compileContract ccip/capability/CCIPHome.sol
+compileContract ccip/NonceManager.sol
+compileContract shared/token/ERC677/BurnMintERC677.sol
+compileContract shared/token/ERC20/BurnMintERC20.sol
+
+
+# Pools
+compileContract ccip/pools/LockReleaseTokenPool.sol
+compileContract ccip/pools/BurnMintTokenPool.sol
+compileContract ccip/pools/BurnFromMintTokenPool.sol
+compileContract ccip/pools/BurnWithFromMintTokenPool.sol
+compileContract ccip/pools/TokenPool.sol
+
+
+# Test helpers
+compileContract ccip/test/helpers/BurnMintERC677Helper.sol
+compileContract ccip/test/helpers/MessageHasher.sol
+compileContract ccip/test/helpers/USDCReaderTester.sol
+compileContract ccip/test/helpers/ReportCodec.sol
+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 (error)
+compileContract ccip/test/helpers/CCIPReaderTester.sol
+
+# Encoding Utils
+# compileContract ccip/interfaces/encodingutils/ICCIPEncodingUtils.sol (error)
+
+# Customer contracts
+compileContract ccip/pools/USDC/USDCTokenPool.sol
+
+compileContract tests/MockV3Aggregator.sol

core/gethwrappers/abigen.go

@@ -28,7 +28,7 @@ const headerComment = `// Code generated - DO NOT EDIT.
 // 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 +72,9 @@ func Abigen(a AbigenArgs) {
 	}
 
 	ImproveAbigenOutput(a.Out, a.ABI)
+	if a.ZkBinPath != "" {
+		ImproveAbigenOutput_zks(a.Out, a.ZkBinPath)
+	}
 }
 
 func ImproveAbigenOutput(path string, abiPath string) {
@@ -466,3 +469,197 @@ func writeInterface(contractName string, fileNode *ast.File) *ast.File {
 func addHeader(code []byte) []byte {
 	return utils.ConcatBytes([]byte(headerComment), code)
 }
+
+// ZK stack logic
+func ImproveAbigenOutput_zks(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/ccip/generated_zks")
+
+		// Extract the parameters from the existing function x
+		paramList := getConstructorParams(x.Type.Params.List)
+		// get the `if zksync()` block
+		zkSyncBlock := getZKSyncBlock(contractName, paramList)
+		// insert the `if zksync()` block
+		addZKSyncBlock(*x, zkSyncBlock)
+		// update the return type in the function signature
+		updateTxReturnType(*x)
+		// update the actual return value
+		updateReturnStmt(*x)
+
+		return false
+	}, nil).(*ast.File)
+}
+
+// get the `if zksync()` block
+func getZKSyncBlock(contractName, paramList 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 = strings.ReplaceAll(zkSyncBlock, "%s", contractName)
+	zkSyncBlock = strings.ReplaceAll(zkSyncBlock, "%params", 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
+}
+
+// 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
+	}
+}

core/gethwrappers/ccip/gen_to_gen_zks.py

@@ -0,0 +1,27 @@
+# utility to convert go_generate.go to go_generate_zks.go
+
+def update_go_generate(file_path):
+    updated_lines = []
+
+    with open(file_path, 'r') as file:
+        for line in file:
+            if "//go:generate" in line:
+                parts = line.strip().split()  # Split the line into parts
+                bin_path = parts[5]
+                # Generate the .zbin path with `.sol` insertion
+                zbin_path = bin_path.replace("solc", "zksolc").replace("/v0.8.24/", "/v1.5.6/").replace(".bin", ".sol/") + bin_path.split('/')[-1].replace(".bin", ".zbin")
+                parts.append(zbin_path)
+
+                parts[3] = "./generation/generate_zks/wrap.go"
+
+                line = " ".join(parts) + "\n"
+                print(line)
+
+            updated_lines.append(line)
+
+    # Write back to the file
+    with open(file_path, 'w') as file:
+        file.writelines(updated_lines)
+
+# Example usage
+update_go_generate('go_generate_zks.go')