github.com/gnolang/gno@v0.0.0-20240520182011-228e9d0192ce/gnovm/pkg/repl/repl.go

package repl

import (
	"bufio"
	"bytes"
	"fmt"
	"go/ast"
	"go/format"
	"go/parser"
	"go/printer"
	"go/token"
	"io"
	"text/template"

	gno "github.com/gnolang/gno/gnovm/pkg/gnolang"
	"github.com/gnolang/gno/gnovm/tests"
)

const (
	executedFunc = "main"
	fileTemplate = `// generated by 'gno repl'
// FILE NAME: test{{$.ID}}.gno
package test
{{range .Imports}}
{{.}}
{{end}}
{{range .Declarations}}
{{.}}
{{end}}
{{with .Expression}}
// main function
func main{{$.ID}}() {
	{{.}}
}
{{end}}
`
)

type tempModel struct {
	ID           int
	Declarations []string
	Imports      []string
	Expression   string
}

type state struct {
	// imports contains all the imports added. they will be added to the following generated source code.
	imports map[string]string

	// files are all generated files. This is used when source code is requested.
	files map[string]string

	// id is the actual execution number. Used to avoid duplicated functions and file names.
	id int

	fileset *token.FileSet
	machine *gno.Machine
}

func newState(stdout io.Writer, sf func() gno.Store) *state {
	s := &state{
		imports: make(map[string]string),
		files:   make(map[string]string),
		fileset: token.NewFileSet(),
	}

	s.machine = gno.NewMachineWithOptions(gno.MachineOptions{
		PkgPath: "test",
		Output:  stdout,
		Store:   sf(),
	})

	return s
}

type ReplOption func(*Repl)

// WithStore allows to modify the default Store implementation used by the VM.
// If nil is provided, the VM will use a default implementation.
func WithStore(s gno.Store) ReplOption {
	return func(r *Repl) {
		r.storeFunc = func() gno.Store {
			return s
		}
	}
}

// WithStd changes std's reader and writers implementations. An internal bytes.Buffer is used by default.
func WithStd(stdin io.Reader, stdout, stderr io.Writer) ReplOption {
	return func(r *Repl) {
		r.stdin = stdin
		r.stdout = stdout
		r.stderr = stderr
	}
}

type Repl struct {
	state *state
	tmpl  *template.Template

	// rw joins stdout and stderr to give an unified output and group with stdin.
	rw bufio.ReadWriter

	// Repl options:
	storeFunc func() gno.Store
	stdout    io.Writer
	stderr    io.Writer
	stdin     io.Reader
}

// NewRepl creates a Repl struct. It is able to process input source code and eventually run it.
func NewRepl(opts ...ReplOption) *Repl {
	t := template.Must(template.New("tmpl").Parse(fileTemplate))

	r := &Repl{
		tmpl: t,
	}

	var b bytes.Buffer
	r.stdin = &b
	r.stdout = &b
	r.stderr = &b

	r.storeFunc = func() gno.Store {
		return tests.TestStore("teststore", "", r.stdin, r.stdout, r.stderr, tests.ImportModeStdlibsOnly)
	}

	for _, o := range opts {
		o(r)
	}

	r.state = newState(r.stdout, r.storeFunc)

	br := bufio.NewReader(r.stdin)
	bw := bufio.NewWriter(io.MultiWriter(r.stderr, r.stdout))
	r.rw = *bufio.NewReadWriter(br, bw)

	return r
}

// Process accepts any valid Gno source code and executes it if it
// is an expression, or stores it for later use if they are declarations.
// If the provided input is not valid Gno source code, an error is returned.
// If the execution on the VM is not successful, the panic is recovered and
// returned as an error.
func (r *Repl) Process(input string) (out string, err error) {
	defer func() {
		if r := recover(); r != nil {
			err = fmt.Errorf("recovered from panic: %q", r)
		}
	}()
	r.state.id++

	decl, declErr := r.parseDeclaration(input)
	if declErr == nil {
		return r.handleDeclarations(decl)
	}

	exp, expErr := r.parseExpression(input)
	if expErr == nil {
		return r.handleExpression(exp)
	}

	return "", fmt.Errorf("error parsing code:\n\t- as expression: %w\n\t- as declarations: %w", expErr, declErr)
}

func (r *Repl) handleExpression(e *ast.File) (string, error) {
	fn := r.filename()
	src := r.nodeToString(e)

	n := gno.MustParseFile(fn, src)
	r.state.files[fn] = src
	r.state.machine.RunFiles(n)
	r.state.machine.RunStatement(gno.S(gno.Call(gno.X(fmt.Sprintf("%s%d", executedFunc, r.state.id)))))

	// Read the result from the output buffer after calling main function.
	b, err := io.ReadAll(r.rw)
	if err != nil {
		return "", fmt.Errorf("error reading output buffer: %w", err)
	}

	return string(b), nil
}

func (r *Repl) handleDeclarations(fn *ast.File) (string, error) {
	var (
		b               bytes.Buffer
		nonImportsCount int
	)
	ast.Inspect(fn, func(n ast.Node) bool {
		var (
			writeNode bool
			ns        string
		)

		switch t := n.(type) {
		case *ast.GenDecl:
			tok := t.Tok

			if tok != token.IMPORT &&
				tok != token.TYPE &&
				tok != token.CONST &&
				tok != token.VAR {
				break
			}

			writeNode = true
			ns = r.nodeToString(n)

			if tok != token.IMPORT {
				nonImportsCount++
				break
			}

			i, ok := t.Specs[0].(*ast.ImportSpec)
			if !ok {
				break
			}

			r.state.imports[i.Path.Value] = ns
		case *ast.FuncDecl:
			writeNode = true
			nonImportsCount++
			ns = r.nodeToString(n)
		}

		if !writeNode {
			return true
		}

		b.WriteString(ns)
		b.WriteByte('\n')

		return false
	})

	// Avoid adding files with only imports on it.
	if nonImportsCount != 0 {
		name := r.filename()
		src := r.nodeToString(fn)

		n := gno.MustParseFile(name, src)
		r.state.files[name] = src
		r.state.machine.RunFiles(n)
	}

	return b.String(), nil
}

func (r *Repl) nodeToString(n ast.Node) string {
	var b bytes.Buffer
	if err := printer.Fprint(&b, r.state.fileset, n); err != nil {
		panic(err)
	}
	return b.String()
}

func (r *Repl) filename() string {
	return fmt.Sprintf("test%d.gno", r.state.id)
}

func (r *Repl) parseExpression(input string) (*ast.File, error) {
	return r.executeAndParse(&tempModel{
		Expression: input,
	})
}

func (r *Repl) parseDeclaration(input string) (*ast.File, error) {
	return r.executeAndParse(&tempModel{
		Declarations: []string{input},
	})
}

func (r *Repl) executeAndParse(m *tempModel) (*ast.File, error) {
	var b bytes.Buffer

	// Set the id to the template.
	m.ID = r.state.id

	// Add all the imports to all the files, they will be removed after formatting.
	for _, e := range r.state.imports {
		m.Imports = append(m.Imports, e)
	}

	err := r.tmpl.Execute(&b, m)
	if err != nil {
		return nil, fmt.Errorf("error executing the template: %w", err)
	}

	// Format the source code to remove unused imports and improve code readability.
	formatted, err := format.Source(b.Bytes())
	if err != nil {
		return nil, fmt.Errorf("error formatting code: %w", err)
	}

	name := r.filename()
	return parser.ParseFile(r.state.fileset, name, formatted, parser.AllErrors|parser.ParseComments)
}

// Reset will reset the actual repl state, restarting the internal VM.
func (r *Repl) Reset() {
	r.state.machine.Release()
	r.state = newState(r.stdout, r.storeFunc)
}

const separator = "//----------------------------------//\n"

// Src will print all the valid code introduced on this Repl session.
func (r *Repl) Src() string {
	var b bytes.Buffer

	b.WriteString(separator)
	for _, s := range r.state.files {
		b.WriteString(s)
		b.WriteString(separator)
	}

	return b.String()
}