Add IR for calls/queries.

aviate-labs · May 3, 2024 · 01381a6 · 01381a6
1 parent 3eb62fc
commit 01381a6
Show file tree

Hide file tree

Showing 27 changed files with 894 additions and 106 deletions.
diff --git a/Makefile b/Makefile
@@ -11,6 +11,7 @@ test-cover:
 	go tool cover -html=coverage.out
 
 gen:
+	cd pocketic && go generate
 	cd candid && go generate
 
 gen-ic:

diff --git a/agent.go b/agent.go
@@ -137,17 +137,26 @@ func New(cfg Config) (*Agent, error) {
 
 // Call calls a method on a canister and unmarshals the result into the given values.
 func (a Agent) Call(canisterID principal.Principal, methodName string, args []any, values []any) error {
-	rawArgs, err := idl.Marshal(args)
+	call, err := a.CreateCall(canisterID, methodName, args...)
 	if err != nil {
 		return err
 	}
+	return call.CallAndWait(values...)
+}
+
+// CreateCall creates a new call to the given canister and method.
+func (a *Agent) CreateCall(canisterID principal.Principal, methodName string, args ...any) (*Call, error) {
+	rawArgs, err := idl.Marshal(args)
+	if err != nil {
+		return nil, err
+	}
 	if len(args) == 0 {
 		// Default to the empty Candid argument list.
 		rawArgs = []byte{'D', 'I', 'D', 'L', 0, 0}
 	}
 	nonce, err := newNonce()
 	if err != nil {
-		return err
+		return nil, err
 	}
 	requestID, data, err := a.sign(Request{
 		Type:          RequestTypeCall,
@@ -159,19 +168,49 @@ func (a Agent) Call(canisterID principal.Principal, methodName string, args []an
 		Nonce:         nonce,
 	})
 	if err != nil {
-		return err
-	}
-	ecID := effectiveCanisterID(canisterID, args)
-	a.logger.Printf("[AGENT] CALL %s %s (%x)", ecID, methodName, *requestID)
-	if _, err := a.call(ecID, data); err != nil {
-		return err
+		return nil, err
 	}
+	return &Call{
+		a:                   a,
+		methodName:          methodName,
+		effectiveCanisterID: effectiveCanisterID(canisterID, args),
+		requestID:           *requestID,
+		data:                data,
+	}, nil
+}
 
-	raw, err := a.poll(ecID, *requestID)
+// CreateQuery creates a new query to the given canister and method.
+func (a *Agent) CreateQuery(canisterID principal.Principal, methodName string, args ...any) (*Query, error) {
+	rawArgs, err := idl.Marshal(args)
 	if err != nil {
-		return err
+		return nil, err
 	}
-	return idl.Unmarshal(raw, values)
+	if len(args) == 0 {
+		// Default to the empty Candid argument list.
+		rawArgs = []byte{'D', 'I', 'D', 'L', 0, 0}
+	}
+	nonce, err := newNonce()
+	if err != nil {
+		return nil, err
+	}
+	_, data, err := a.sign(Request{
+		Type:          RequestTypeQuery,
+		Sender:        a.Sender(),
+		CanisterID:    canisterID,
+		MethodName:    methodName,
+		Arguments:     rawArgs,
+		IngressExpiry: a.expiryDate(),
+		Nonce:         nonce,
+	})
+	if err != nil {
+		return nil, err
+	}
+	return &Query{
+		a:                   a,
+		methodName:          methodName,
+		effectiveCanisterID: effectiveCanisterID(canisterID, args),
+		data:                data,
+	}, nil
 }
 
 // GetCanisterControllers returns the list of principals that can control the given canister.
@@ -245,51 +284,18 @@ func (a Agent) GetCanisterModuleHash(canisterID principal.Principal) ([]byte, er
 	return h, err
 }
 
+// GetRootKey returns the root key of the host.
 func (a Agent) GetRootKey() []byte {
 	return a.rootKey
 }
 
+// Query calls a method on a canister and unmarshals the result into the given values.
 func (a Agent) Query(canisterID principal.Principal, methodName string, args []any, values []any) error {
-	rawArgs, err := idl.Marshal(args)
+	query, err := a.CreateQuery(canisterID, methodName, args...)
 	if err != nil {
 		return err
 	}
-	if len(args) == 0 {
-		// Default to the empty Candid argument list.
-		rawArgs = []byte{'D', 'I', 'D', 'L', 0, 0}
-	}
-	nonce, err := newNonce()
-	if err != nil {
-		return err
-	}
-	_, data, err := a.sign(Request{
-		Type:          RequestTypeQuery,
-		Sender:        a.Sender(),
-		CanisterID:    canisterID,
-		MethodName:    methodName,
-		Arguments:     rawArgs,
-		IngressExpiry: a.expiryDate(),
-		Nonce:         nonce,
-	})
-	if err != nil {
-		return err
-	}
-	ecID := effectiveCanisterID(canisterID, args)
-	a.logger.Printf("[AGENT] QUERY %s %s", ecID, methodName)
-	resp, err := a.query(ecID, data)
-	if err != nil {
-		return err
-	}
-	var raw []byte
-	switch resp.Status {
-	case "replied":
-		raw = resp.Reply["arg"]
-	case "rejected":
-		return fmt.Errorf("(%d) %s", resp.RejectCode, resp.RejectMsg)
-	default:
-		panic("unreachable")
-	}
-	return idl.Unmarshal(raw, values)
+	return query.Query(values...)
 }
 
 // RequestStatus returns the status of the request with the given ID.
@@ -332,8 +338,8 @@ func (a Agent) Sender() principal.Principal {
 	return a.identity.Sender()
 }
 
-func (a Agent) call(ecid principal.Principal, data []byte) ([]byte, error) {
-	return a.client.call(ecid, data)
+func (a Agent) call(ecID principal.Principal, data []byte) ([]byte, error) {
+	return a.client.call(ecID, data)
 }
 
 func (a Agent) expiryDate() uint64 {
@@ -428,6 +434,34 @@ func (a Agent) sign(request Request) (*RequestID, []byte, error) {
 	return &requestID, data, nil
 }
 
+// Call is an intermediate representation of a call to a canister.
+type Call struct {
+	a                   *Agent
+	methodName          string
+	effectiveCanisterID principal.Principal
+	requestID           RequestID
+	data                []byte
+}
+
+// CallAndWait calls a method on a canister and waits for the result.
+func (c Call) CallAndWait(values ...any) error {
+	c.a.logger.Printf("[AGENT] CALL %s %s (%x)", c.effectiveCanisterID, c.methodName, c.requestID)
+	if _, err := c.a.call(c.effectiveCanisterID, c.data); err != nil {
+		return err
+	}
+	raw, err := c.a.poll(c.effectiveCanisterID, c.requestID)
+	if err != nil {
+		return err
+	}
+	return idl.Unmarshal(raw, values)
+}
+
+// WithEffectiveCanisterID sets the effective canister ID for the call.
+func (c *Call) WithEffectiveCanisterID(canisterID principal.Principal) *Call {
+	c.effectiveCanisterID = canisterID
+	return c
+}
+
 // Config is the configuration for an Agent.
 type Config struct {
 	// Identity is the identity used by the Agent.
@@ -445,3 +479,30 @@ type Config struct {
 	// PollTimeout is the timeout for polling for a response.
 	PollTimeout time.Duration
 }
+
+// Query is an intermediate representation of a query to a canister.
+type Query struct {
+	a                   *Agent
+	methodName          string
+	effectiveCanisterID principal.Principal
+	data                []byte
+}
+
+// Query calls a method on a canister and unmarshals the result into the given values.
+func (q Query) Query(values ...any) error {
+	q.a.logger.Printf("[AGENT] QUERY %s %s", q.effectiveCanisterID, q.methodName)
+	resp, err := q.a.query(q.effectiveCanisterID, q.data)
+	if err != nil {
+		return err
+	}
+	var raw []byte
+	switch resp.Status {
+	case "replied":
+		raw = resp.Reply["arg"]
+	case "rejected":
+		return fmt.Errorf("(%d) %s", resp.RejectCode, resp.RejectMsg)
+	default:
+		panic("unreachable")
+	}
+	return idl.Unmarshal(raw, values)
+}
diff --git a/candid/internal/candid/grammar.go b/candid/internal/candid/grammar.go
@@ -643,7 +643,9 @@ func Record(p *ast.Parser) (*ast.Node, error) {
 			TypeStrings: NodeTypes,
 			Value: op.And{
 				"record",
-				Sp,
+				op.Optional(
+					Sp,
+				),
 				'{',
 				Ws,
 				op.Optional(

diff --git a/candid/internal/candid/grammar.pegn b/candid/internal/candid/grammar.pegn
@@ -31,7 +31,7 @@ ConsType   <- Blob
 Blob      <-- 'blob'
 Opt       <-- 'opt' Sp DataType
 Vec       <-- 'vec' Sp DataType
-Record    <-- 'record' Sp '{' Ws Fields? Ws '}'
+Record    <-- 'record' Sp? '{' Ws Fields? Ws '}'
 Variant   <-- 'variant' Sp '{' Ws Fields? Ws '}'
 Fields     <- FieldType (';' Ws FieldType)* ';'?
 

diff --git a/didc b/didc
diff --git a/gen/generator.go b/gen/generator.go
@@ -76,6 +76,8 @@ type Generator struct {
 	PackageName        string
 	ServiceDescription did.Description
 	usedIDL            bool
+
+	indirect bool
 }
 
 // NewGenerator creates a new generator for the given service description.
@@ -145,7 +147,11 @@ func (g *Generator) Generate() ([]byte, error) {
 			})
 		}
 	}
-	t, ok := templates["agent"]
+	tmplName := "agent"
+	if g.indirect {
+		tmplName = "agent_indirect"
+	}
+	t, ok := templates[tmplName]
 	if !ok {
 		return nil, fmt.Errorf("template not found")
 	}
@@ -320,6 +326,12 @@ func (g *Generator) GenerateMock() ([]byte, error) {
 	return io.ReadAll(&tmpl)
 }
 
+// Indirect sets the generator to generate indirect calls.
+func (g *Generator) Indirect() *Generator {
+	g.indirect = true
+	return g
+}
+
 func (g *Generator) dataToGoReturnValue(definitions map[string]did.Data, prefix string, data did.Data) string {
 	switch t := data.(type) {
 	case did.Primitive:

diff --git a/gen/templates/agent_indirect.gotmpl b/gen/templates/agent_indirect.gotmpl
@@ -0,0 +1,60 @@
+// Package {{ .PackageName }} provides a client for the "{{ .CanisterName }}" canister.
+// Do NOT edit this file. It was automatically generated by https://github.com/aviate-labs/agent-go.
+package {{ .PackageName }}
+
+import (
+    "github.com/aviate-labs/agent-go"
+    {{ if .UsedIDL }}"github.com/aviate-labs/agent-go/candid/idl"{{ end }}
+    "github.com/aviate-labs/agent-go/principal"
+)
+
+{{- range .Definitions }}
+
+type {{ .Name }} {{ if .Eq }}= {{end}}{{ .Type }}
+{{- end }}
+
+// {{ .AgentName }}Agent is a client for the "{{ .CanisterName }}" canister.
+type {{ .AgentName }}Agent struct {
+    a          *agent.Agent
+    canisterId principal.Principal
+}
+
+// New{{ .AgentName }}Agent creates a new agent for the "{{ .CanisterName }}" canister.
+func New{{ .AgentName }}Agent(canisterId principal.Principal, config agent.Config) (*{{ .AgentName }}Agent, error) {
+    a, err := agent.New(config)
+    if err != nil {
+        return nil, err
+    }
+    return &{{ .AgentName }}Agent{
+        a:          a,
+        canisterId: canisterId,
+    }, nil
+}
+{{- range .Methods }}
+
+// {{ .Name }} calls the "{{ .RawName }}" method on the "{{ $.CanisterName }}" canister.
+func (a {{ $.AgentName }}Agent) {{ .Name }}({{ range $i, $e := .ArgumentTypes }}{{ if $i }}, {{ end }}{{ $e.Name }} {{ $e.Type }}{{ end }}) {{ if .ReturnTypes }}({{ range .ReturnTypes }}*{{ . }}, {{ end }}error){{ else }}error{{ end }} {
+    {{ range $i, $e := .ReturnTypes -}}
+        var r{{ $i }} {{ $e }}
+    {{ end -}}
+    if err := a.a.{{ .Type }}(
+        a.canisterId,
+        "{{ .RawName }}",
+        []any{{ "{" }}{{ range $i, $e := .ArgumentTypes }}{{ if $i }}, {{ end }}{{ $e.Name }}{{ end }}{{ "}" }},
+        []any{{ "{" }}{{ range $i, $e := .ReturnTypes }}{{ if $i }}, {{ end }}&r{{ $i }}{{ end }}{{ "}"}},
+    ); err != nil {
+        return {{ range .ReturnTypes }}nil, {{ end }}err
+    }
+    return {{ range $i, $_ := .ReturnTypes }}&r{{ $i }}, {{ end }}nil
+}
+
+// {{ .Name }}{{ .Type }} creates an indirect representation of the "{{ .RawName }}" method on the "{{ $.CanisterName }}" canister.
+func (a {{ $.AgentName }}Agent) {{ .Name }}{{ .Type }}({{ range $i, $e := .ArgumentTypes }}{{ if $i }}, {{ end }}{{ $e.Name }} {{ $e.Type }}{{ end }}) (*agent.{{ .Type }},error) {
+    return a.a.Create{{ .Type }}(
+        a.canisterId,
+        "{{ .RawName }}",{{ range $i, $e := .ArgumentTypes }}
+        {{ $e.Name }},{{ end }}
+    )
+}
+
+{{- end }}