[x/programs] Add host abstractions

x/programs/examples/counter_test.go

@@ -17,6 +17,7 @@ import (
 	"github.com/ava-labs/hypersdk/x/programs/examples/imports/program"
 	"github.com/ava-labs/hypersdk/x/programs/examples/imports/pstate"
 	"github.com/ava-labs/hypersdk/x/programs/examples/storage"
+	"github.com/ava-labs/hypersdk/x/programs/host"
 	"github.com/ava-labs/hypersdk/x/programs/runtime"
 	"github.com/ava-labs/hypersdk/x/programs/tests"
 )
@@ -40,16 +41,17 @@ func TestCounterProgram(t *testing.T) {
 
 	eng := engine.New(engine.NewConfig())
 	// define supported imports
-	supported := runtime.NewSupportedImports()
-	supported.Register("state", func() runtime.Import {
+	importsBuilder := host.NewImportsBuilder()
+	importsBuilder.Register("state", func() host.Import {
 		return pstate.New(log, db)
 	})
-	supported.Register("program", func() runtime.Import {
+	importsBuilder.Register("program", func() host.Import {
 		return program.New(log, eng, db, cfg)
 	})
+	imports := importsBuilder.Build()
 
 	wasmBytes := tests.ReadFixture(t, "../tests/fixture/counter.wasm")
-	rt := runtime.New(log, eng, supported.Imports(), cfg)
+	rt := runtime.New(log, eng, imports, cfg)
 	err := rt.Initialize(ctx, wasmBytes, maxUnits)
 	require.NoError(err)
 
@@ -85,7 +87,7 @@ func TestCounterProgram(t *testing.T) {
 
 	// initialize second runtime to create second counter program with an empty
 	// meter.
-	rt2 := runtime.New(log, eng, supported.Imports(), cfg)
+	rt2 := runtime.New(log, eng, imports, cfg)
 	err = rt2.Initialize(ctx, wasmBytes, engine.NoUnits)
 
 	require.NoError(err)

x/programs/examples/imports/program/program.go

@@ -6,7 +6,6 @@ package program
 import (
 	"context"
 	"encoding/binary"
-	"fmt"
 
 	"go.uber.org/zap"
 
@@ -19,26 +18,29 @@ import (
 
 	"github.com/ava-labs/hypersdk/x/programs/engine"
 	"github.com/ava-labs/hypersdk/x/programs/examples/storage"
+	"github.com/ava-labs/hypersdk/x/programs/host"
 	"github.com/ava-labs/hypersdk/x/programs/runtime"
 )
 
+var _ host.Import = (*Import)(nil)
+
 const Name = "program"
 
 type Import struct {
-	cfg        *runtime.Config
-	db         state.Mutable
-	log        logging.Logger
-	imports    runtime.SupportedImports
-	meter      *engine.Meter
-	engine     *engine.Engine
-	registered bool
+	mu  state.Mutable
+	log logging.Logger
+	cfg *runtime.Config
+
+	engine  *engine.Engine
+	meter   *engine.Meter
+	imports host.SupportedImports
 }
 
 // New returns a new program invoke host module which can perform program to program calls.
-func New(log logging.Logger, engine *engine.Engine, db state.Mutable, cfg *runtime.Config) *Import {
+func New(log logging.Logger, engine *engine.Engine, mu state.Mutable, cfg *runtime.Config) *Import {
 	return &Import{
 		cfg:    cfg,
-		db:     db,
+		mu:     mu,
 		log:    log,
 		engine: engine,
 	}
@@ -48,18 +50,10 @@ func (i *Import) Name() string {
 	return Name
 }
 
-func (i *Import) Register(link runtime.Link, meter *engine.Meter, imports runtime.SupportedImports) error {
-	if i.registered {
-		return fmt.Errorf("import module already registered: %q", Name)
-	}
-	i.imports = imports
-	i.meter = meter
-
-	if err := link.FuncWrap(Name, "call_program", i.callProgramFn); err != nil {
-		return err
-	}
-
-	return nil
+func (i *Import) Register(link *host.Link) error {
+	i.meter = link.Meter()
+	i.imports = link.Imports()
+	return link.RegisterImportFn(Name, "call_program", i.callProgramFn)
 }
 
 // callProgramFn makes a call to an entry function of a program in the context of another program's ID.
@@ -73,6 +67,7 @@ func (i *Import) callProgramFn(
 ) int64 {
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
+
 	memory := runtime.NewMemory(runtime.NewExportClient(caller))
 	// get the entry function for invoke to call.
 	functionBytes, err := runtime.SmartPtr(function).Bytes(memory)
@@ -92,7 +87,7 @@ func (i *Import) callProgramFn(
 	}
 
 	// get the program bytes from storage
-	programWasmBytes, err := getProgramWasmBytes(i.log, i.db, programIDBytes)
+	programWasmBytes, err := getProgramWasmBytes(i.log, i.mu, programIDBytes)
 	if err != nil {
 		i.log.Error("failed to get program bytes from storage",
 			zap.Error(err),

x/programs/examples/imports/pstate/pstate.go

@@ -6,7 +6,6 @@ package pstate
 import (
 	"context"
 	"errors"
-	"fmt"
 
 	"github.com/bytecodealliance/wasmtime-go/v14"
 
@@ -18,45 +17,36 @@ import (
 	"github.com/ava-labs/hypersdk/state"
 	"github.com/ava-labs/hypersdk/x/programs/engine"
 	"github.com/ava-labs/hypersdk/x/programs/examples/storage"
+	"github.com/ava-labs/hypersdk/x/programs/host"
 	"github.com/ava-labs/hypersdk/x/programs/runtime"
 )
 
-const Name = "state"
+var _ host.Import = (*Import)(nil)
 
-var _ runtime.Import = &Import{}
+const Name = "state"
 
 // New returns a program storage module capable of storing arbitrary bytes
 // in the program's namespace.
-func New(log logging.Logger, mu state.Mutable) runtime.Import {
+func New(log logging.Logger, mu state.Mutable) host.Import {
 	return &Import{mu: mu, log: log}
 }
 
 type Import struct {
-	mu         state.Mutable
-	log        logging.Logger
-	meter      *engine.Meter
-	registered bool
+	mu    state.Mutable
+	log   logging.Logger
+	meter *engine.Meter
 }
 
 func (i *Import) Name() string {
 	return Name
 }
 
-func (i *Import) Register(link runtime.Link, meter *engine.Meter, _ runtime.SupportedImports) error {
-	if i.registered {
-		return fmt.Errorf("import module already registered: %q", Name)
-	}
-	i.meter = meter
-	i.registered = true
-
-	if err := link.FuncWrap(Name, "put", i.putFn); err != nil {
+func (i *Import) Register(link *host.Link) error {
+	i.meter = link.Meter()
+	if err := link.RegisterImportFn(Name, "put", i.putFn); err != nil {
 		return err
 	}
-	if err := link.FuncWrap(Name, "get", i.getFn); err != nil {
-		return err
-	}
-
-	return nil
+	return link.RegisterImportFn(Name, "get", i.getFn)
 }
 
 func (i *Import) putFn(caller *wasmtime.Caller, id int64, key int64, value int64) int32 {

x/programs/examples/token.go

@@ -16,10 +16,11 @@ import (
 	"github.com/ava-labs/hypersdk/state"
 	"github.com/ava-labs/hypersdk/x/programs/engine"
 	"github.com/ava-labs/hypersdk/x/programs/examples/storage"
+	"github.com/ava-labs/hypersdk/x/programs/host"
 	"github.com/ava-labs/hypersdk/x/programs/runtime"
 )
 
-func NewToken(log logging.Logger, engine *engine.Engine, programBytes []byte, db state.Mutable, cfg *runtime.Config, imports runtime.SupportedImports, maxUnits uint64) *Token {
+func NewToken(log logging.Logger, engine *engine.Engine, programBytes []byte, db state.Mutable, cfg *runtime.Config, imports host.SupportedImports, maxUnits uint64) *Token {
 	return &Token{
 		log:          log,
 		programBytes: programBytes,
@@ -42,7 +43,7 @@ type Token struct {
 	log          logging.Logger
 	programBytes []byte
 	cfg          *runtime.Config
-	imports      runtime.SupportedImports
+	imports      host.SupportedImports
 	db           state.Mutable
 	maxUnits     uint64
 	engine       *engine.Engine

x/programs/examples/token_test.go

@@ -14,6 +14,7 @@ import (
 
 	"github.com/ava-labs/hypersdk/x/programs/engine"
 	"github.com/ava-labs/hypersdk/x/programs/examples/imports/pstate"
+	"github.com/ava-labs/hypersdk/x/programs/host"
 	"github.com/ava-labs/hypersdk/x/programs/runtime"
 	"github.com/ava-labs/hypersdk/x/programs/tests"
 )
@@ -113,10 +114,10 @@ func newTokenProgram(maxUnits uint64, engine *engine.Engine, cfg *runtime.Config
 		))
 
 	// define imports
-	supported := runtime.NewSupportedImports()
-	supported.Register("state", func() runtime.Import {
+	importsBuilder := host.NewImportsBuilder()
+	importsBuilder.Register("state", func() host.Import {
 		return pstate.New(log, db)
 	})
 
-	return NewToken(log, engine, programBytes, db, cfg, supported.Imports(), maxUnits), nil
+	return NewToken(log, engine, programBytes, db, cfg, importsBuilder.Build(), maxUnits), nil
 }

x/programs/host/dependencies.go

@@ -0,0 +1,11 @@
+// Copyright (C) 2023, Ava Labs, Inc. All rights reserved.
+// See the file LICENSE for licensing terms.
+
+package host
+
+type Import interface {
+	// Name returns the name of this import module.
+	Name() string
+	// Register registers this import module with the provided link.
+	Register(*Link) error
+}

x/programs/host/imports.go

@@ -0,0 +1,67 @@
+// Copyright (C) 2023, Ava Labs, Inc. All rights reserved.
+// See the file LICENSE for licensing terms.
+
+package host
+
+import (
+	"github.com/bytecodealliance/wasmtime-go/v14"
+
+	"github.com/ava-labs/hypersdk/x/programs/engine"
+)
+
+const (
+	wasiPreview1ModName = "wasi_snapshot_preview1"
+)
+
+var NoSupportedImports = make(SupportedImports)
+
+type ImportFnCallback struct {
+	// beforeRequest is called before the import function request is made.
+	BeforeRequest func(module, name string, meter *engine.Meter) error
+	// afterResponse is called after the import function response is received.
+	AfterResponse func(module, name string, meter *engine.Meter) error
+}
+
+// Supported is a map of supported import modules. The runtime will enable these imports
+// during initialization only if implemented by the `program`.
+type SupportedImports map[string]func() Import
+
+type ImportsBuilder struct {
+	imports map[string]func() Import
+}
+
+func NewImportsBuilder() *ImportsBuilder {
+	return &ImportsBuilder{
+		imports: make(map[string]func() Import),
+	}
+}
+
+// Register registers a supported import by name.
+func (s *ImportsBuilder) Register(name string, f func() Import) *ImportsBuilder {
+	s.imports[name] = f
+	return s
+}
+
+// Build returns the supported imports.
+func (s *ImportsBuilder) Build() SupportedImports {
+	return s.imports
+}
+
+// getRegisteredImports returns the unique names of all import modules registered
+// by the wasm module.
+func getRegisteredImports(importTypes []*wasmtime.ImportType) []string {
+	u := make(map[string]struct{}, len(importTypes))
+	imports := []string{}
+	for _, t := range importTypes {
+		mod := t.Module()
+		if mod == wasiPreview1ModName {
+			continue
+		}
+		if _, ok := u[mod]; ok {
+			continue
+		}
+		u[mod] = struct{}{}
+		imports = append(imports, mod)
+	}
+	return imports
+}