🪚 OmniGraph™ Transformations & schemas (#76)

packages/test-utils/src/arbitraries.ts

@@ -2,6 +2,10 @@ import fc from 'fast-check'
 import { EndpointId, Stage } from '@layerzerolabs/lz-definitions'
 import { ENDPOINT_IDS } from './constants'
 
+export const nullishArbitrary = fc.constantFrom(null, undefined)
+
+export const nullableArbitrary = <T>(a: fc.Arbitrary<T>) => fc.oneof(a, nullishArbitrary)
+
 export const addressArbitrary = fc.string()
 
 export const evmAddressArbitrary = fc.hexaString({ minLength: 40, maxLength: 40 }).map((address) => `0x${address}`)

packages/ua-utils-evm-hardhat-test/test/__utils__/endpoint.ts

@@ -144,11 +144,11 @@ export const setupDefaultEndpoint = async (): Promise<void> => {
     }
 
     // Now we compile a list of all the transactions that need to be executed for the ULNs and Endpoints
-    const builderEndpoint = await OmniGraphBuilderHardhat.fromConfig(config, contractFactory)
+    const builderEndpoint = await OmniGraphBuilderHardhat.fromConfig(config)
     const endpointTransactions = await configureEndpoint(builderEndpoint.graph, endpointSdkFactory)
-    const builderSendUln = await OmniGraphBuilderHardhat.fromConfig(sendUlnConfig, contractFactory)
+    const builderSendUln = await OmniGraphBuilderHardhat.fromConfig(sendUlnConfig)
     const sendUlnTransactions = await configureUln302(builderSendUln.graph, ulnSdkFactory)
-    const builderReceiveUln = await OmniGraphBuilderHardhat.fromConfig(receiveUlnConfig, contractFactory)
+    const builderReceiveUln = await OmniGraphBuilderHardhat.fromConfig(receiveUlnConfig)
     const receiveUlnTransactions = await configureUln302(builderReceiveUln.graph, ulnSdkFactory)
 
     const transactions = [...sendUlnTransactions, ...receiveUlnTransactions, ...endpointTransactions]

packages/ua-utils-evm-hardhat-test/test/oapp/config.test.ts

@@ -54,7 +54,7 @@ describe('oapp/config', () => {
         // This is the required tooling we need to set up
         const contractFactory = createContractFactory()
         const connectedContractFactory = createConnectedContractFactory(contractFactory)
-        const builder = await OmniGraphBuilderHardhat.fromConfig(config, contractFactory)
+        const builder = await OmniGraphBuilderHardhat.fromConfig(config)
 
         // This so far the only non-oneliner, a function that returns an SDK for a contract on a network
         const sdkFactory = async (point: OmniPoint) => new OApp(await connectedContractFactory(point))
@@ -79,7 +79,7 @@ describe('oapp/config', () => {
         // This is the required tooling we need to set up
         const contractFactory = createContractFactory()
         const connectedContractFactory = createConnectedContractFactory(contractFactory)
-        const builder = await OmniGraphBuilderHardhat.fromConfig(config, contractFactory)
+        const builder = await OmniGraphBuilderHardhat.fromConfig(config)
 
         // This so far the only non-oneliner, a function that returns an SDK for a contract on a network
         const sdkFactory = async (point: OmniPoint) => new OApp(await connectedContractFactory(point))

packages/ua-utils-evm-hardhat/src/oapp/index.ts

@@ -0,0 +1 @@
+export * from './types'

packages/ua-utils-evm-hardhat/src/oapp/types.ts

@@ -0,0 +1,3 @@
+import { OmniGraphHardhat } from '@layerzerolabs/utils-evm-hardhat'
+
+export type OAppOmniGraphHardhat = OmniGraphHardhat<unknown, unknown>

packages/utils-evm-hardhat/src/omnigraph/builder.ts

@@ -1,8 +1,7 @@
-import type { OmniEdge, OmniNode } from '@layerzerolabs/utils'
-import type { OmniContractFactoryHardhat, OmniGraphHardhat } from './types'
+import type { OmniGraphHardhat } from './types'
 import { OmniGraphBuilder } from '@layerzerolabs/utils'
-import { omniContractToPoint } from '@layerzerolabs/utils-evm'
 import assert from 'assert'
+import { OmniGraphHardhatTransformer, createOmniGraphHardhatTransformer } from './transformations'
 
 /**
  * OmniGraphBuilderHardhat houses all hardhat-specific utilities for building OmniGraphs
@@ -12,28 +11,9 @@ import assert from 'assert'
 export class OmniGraphBuilderHardhat {
     static async fromConfig<TNodeConfig, TEdgeConfig>(
         graph: OmniGraphHardhat<TNodeConfig, TEdgeConfig>,
-        contractFactory: OmniContractFactoryHardhat
+        transform: OmniGraphHardhatTransformer<TNodeConfig, TEdgeConfig> = createOmniGraphHardhatTransformer()
     ): Promise<OmniGraphBuilder<TNodeConfig, TEdgeConfig>> {
-        const builder = new OmniGraphBuilder<TNodeConfig, TEdgeConfig>()
-
-        const nodes: OmniNode<TNodeConfig>[] = await Promise.all(
-            graph.contracts.map(async ({ contract, config }) => ({
-                point: omniContractToPoint(await contractFactory(contract)),
-                config,
-            }))
-        )
-
-        const edges: OmniEdge<TEdgeConfig>[] = await Promise.all(
-            graph.connections.map(async ({ from, to, config }) => ({
-                vector: {
-                    from: omniContractToPoint(await contractFactory(from)),
-                    to: omniContractToPoint(await contractFactory(to)),
-                },
-                config,
-            }))
-        )
-
-        return builder.addNodes(...nodes).addEdges(...edges)
+        return OmniGraphBuilder.fromGraph(await transform(graph))
     }
 
     constructor() {

packages/utils-evm-hardhat/src/omnigraph/index.ts

@@ -1,4 +1,6 @@
 export * from './builder'
 export * from './contracts'
 export * from './coordinates'
+export * from './schema'
+export * from './transformations'
 export * from './types'

packages/utils-evm-hardhat/src/omnigraph/schema.ts

@@ -0,0 +1,59 @@
+import { z } from 'zod'
+import { EndpointIdSchema, OmniPointSchema } from '@layerzerolabs/utils'
+import type { OmniEdgeHardhat, OmniGraphHardhat, OmniNodeHardhat, OmniPointHardhat } from './types'
+
+export const OmniPointHardhatSchema: z.ZodSchema<OmniPointHardhat, z.ZodTypeDef, unknown> = z.object({
+    eid: EndpointIdSchema,
+    contractName: z.string().nullish(),
+    address: z.string().nullish(),
+})
+
+const OmniPointOrOmniPointHardhatSchema = z.union([OmniPointHardhatSchema, OmniPointSchema])
+
+/**
+ * Factory for OmniNodeHardhat schemas
+ *
+ * @param configSchema Schema of the config contained in the node
+ *
+ * @returns {z.ZodSchema<OmniNodeHardhat<TConfig>>} schema for a node with the particular config type
+ */
+export const createOmniNodeHardhatSchema = <TConfig = unknown>(
+    configSchema: z.ZodSchema<TConfig, z.ZodTypeDef, unknown>
+): z.ZodSchema<OmniNodeHardhat<TConfig>, z.ZodTypeDef, unknown> =>
+    z.object({
+        contract: OmniPointOrOmniPointHardhatSchema,
+        config: configSchema,
+    }) as z.ZodSchema<OmniNodeHardhat<TConfig>, z.ZodTypeDef, unknown>
+
+/**
+ * Factory for OmniEdgeHardhat schemas
+ *
+ * @param {z.ZodSchema<TConfig>} configSchema Schema of the config contained in the edge
+ *
+ * @returns {z.ZodSchema<OmniEdgeHardhat<TConfig>>} Schema for an edge with the particular config type
+ */
+export const createOmniEdgeHardhatSchema = <TConfig = unknown>(
+    configSchema: z.ZodSchema<TConfig, z.ZodTypeDef, unknown>
+): z.ZodSchema<OmniEdgeHardhat<TConfig>, z.ZodTypeDef, unknown> =>
+    z.object({
+        from: OmniPointOrOmniPointHardhatSchema,
+        to: OmniPointOrOmniPointHardhatSchema,
+        config: configSchema,
+    }) as z.ZodSchema<OmniEdgeHardhat<TConfig>, z.ZodTypeDef, unknown>
+
+/**
+ * Factory for OmniGraphHardhat schemas
+ *
+ * @param {z.ZodSchema<OmniNodeHardhat<TNodeConfig>>} nodeSchema
+ * @param {z.ZodSchema<OmniEdgeHardhat<TEdgeConfig>>} edgeSchema
+ *
+ * @returns {z.ZodSchema<OmniGraphHardhat<TNodeConfig, TEdgeConfig>>}
+ */
+export const createOmniGraphHardhatSchema = <TNodeConfig = unknown, TEdgeConfig = unknown>(
+    nodeSchema: z.ZodSchema<OmniNodeHardhat<TNodeConfig>, z.ZodTypeDef, unknown>,
+    edgeSchema: z.ZodSchema<OmniEdgeHardhat<TEdgeConfig>, z.ZodTypeDef, unknown>
+): z.ZodSchema<OmniGraphHardhat<TNodeConfig, TEdgeConfig>, z.ZodTypeDef, unknown> =>
+    z.object({
+        contracts: z.array(nodeSchema),
+        connections: z.array(edgeSchema),
+    })

packages/utils-evm-hardhat/src/omnigraph/transformations.ts

@@ -0,0 +1,43 @@
+import type { OmniEdge, OmniGraph, OmniNode } from '@layerzerolabs/utils'
+import { isOmniPoint } from '@layerzerolabs/utils'
+import { omniContractToPoint } from '@layerzerolabs/utils-evm'
+import { createContractFactory } from './coordinates'
+import type { OmniContractFactoryHardhat, OmniEdgeHardhat, OmniGraphHardhat, OmniNodeHardhat } from './types'
+
+export const createOmniNodeHardhatTransformer =
+    (contractFactory: OmniContractFactoryHardhat = createContractFactory()) =>
+    async <TNodeConfig = unknown>({
+        contract,
+        config,
+    }: OmniNodeHardhat<TNodeConfig>): Promise<OmniNode<TNodeConfig>> => {
+        const point = isOmniPoint(contract) ? contract : omniContractToPoint(await contractFactory(contract))
+
+        return { point, config }
+    }
+
+export const createOmniEdgeHardhatTransformer =
+    (contractFactory: OmniContractFactoryHardhat = createContractFactory()) =>
+    async <TEdgeConfig = unknown>({
+        from: fromContract,
+        to: toContract,
+        config,
+    }: OmniEdgeHardhat<TEdgeConfig>): Promise<OmniEdge<TEdgeConfig>> => {
+        const from = isOmniPoint(fromContract) ? fromContract : omniContractToPoint(await contractFactory(fromContract))
+        const to = isOmniPoint(toContract) ? toContract : omniContractToPoint(await contractFactory(toContract))
+
+        return { vector: { from, to }, config }
+    }
+
+export type OmniGraphHardhatTransformer<TNodeConfig = unknown, TEdgeConfig = unknown> = (
+    graph: OmniGraphHardhat<TNodeConfig, TEdgeConfig>
+) => Promise<OmniGraph<TNodeConfig, TEdgeConfig>>
+
+export const createOmniGraphHardhatTransformer =
+    <TNodeConfig = unknown, TEdgeConfig = unknown>(
+        nodeTransformer = createOmniNodeHardhatTransformer(),
+        edgeTransformer = createOmniEdgeHardhatTransformer()
+    ): OmniGraphHardhatTransformer<TNodeConfig, TEdgeConfig> =>
+    async (graph) => ({
+        contracts: await Promise.all(graph.contracts.map(nodeTransformer)),
+        connections: await Promise.all(graph.connections.map(edgeTransformer)),
+    })

packages/utils-evm-hardhat/src/omnigraph/types.ts

@@ -4,8 +4,8 @@ import type { OmniContract } from '@layerzerolabs/utils-evm'
 
 export interface OmniPointHardhat {
     eid: EndpointId
-    contractName?: string
-    address?: string
+    contractName?: string | null
+    address?: string | null
 }
 
 export interface OmniNodeHardhat<TNodeConfig> {
@@ -14,8 +14,8 @@ export interface OmniNodeHardhat<TNodeConfig> {
 }
 
 export interface OmniEdgeHardhat<TEdgeConfig> {
-    from: OmniPointHardhat
-    to: OmniPointHardhat
+    from: OmniPointHardhat | OmniPoint
+    to: OmniPointHardhat | OmniPoint
     config: TEdgeConfig
 }