argos/ast/ast.go

package ast

import (
	"fmt"
	"strings"

	"github.com/cevaris/ordered_map"
)

type AST struct {
	Tag Tag
	Val interface{}
}

type Tag int8

const (
	Symbol Tag = iota
	List
	String
	Num
	Bool
	Fn
	Char
	Table
	Prim
	Quoted
)

func (tag Tag) String() string {
	names := []string{
		"Symbol",
		"List",
		"String",
		"Num",
		"Bool",
		"Fn",
		"Char",
		"Table",
		"Prim",
		"Quoted",
	}

	return names[tag]
}

func (ast *AST) Type() *AST {
	names := []string{
		"sym",
		"cons",
		"string",
		"num",
		"bool",
		"fn",
		"char",
		"table",
		"fn",
		"cons",
	}

	if ast.Tag == Quoted {
		return ast.Val.(*AST).Type()
	}

	name := names[ast.Tag]
	if ast.Tag == Num {
		val := ast.Val.(float64)
		if val == float64(int64(val)) {
			name = "int"
		}
	}

	res := AST{Symbol, name}
	return &res
}

func (ast *AST) Pretty() string {
	if ast.Tag == List {
		var agg []string
		for _, elem := range ast.Val.([]*AST) {
			agg = append(agg, elem.Pretty())
		}
		return "(" + strings.Join(agg, " ") + ")"
	} else if ast.Tag == Quoted {
		return "'" + ast.Val.(*AST).Pretty()
	} else if ast.Tag == Fn {
		val := ast.Val.(Func)
		var agg []string
		for _, elem := range val.Params {
			agg = append(agg, elem)
		}
		rest := ""
		if val.HasRest() {
			rest = " . " + *val.Rest
		}
		// TODO: opt
		opt := ""
		iter := val.Opt.IterFunc()
		for kv, ok := iter(); ok; kv, ok = iter() {
			opt += "(o " + kv.Key.(string) + " " + kv.Value.(*AST).Pretty() + ")"
		}
		if val.HasOpt() && (val.HasRest() || len(agg) != 0) {
			opt = " " + opt
		}
		body := val.Body.Pretty()
		return "(fn (" + strings.Join(agg, " ") + rest + opt + ") " + body + ")"
	} else if ast.Tag == Char {
		return "#" + string(ast.Val.(rune))
	} else if ast.Tag == Table {
		var agg []string

		for k, v := range ast.Val.(map[*AST]*AST) {
			agg = append(agg, "("+k.Pretty()+" . "+v.Pretty()+")")
		}

		return "#hash(" + strings.Join(agg, " ") + ")"
	} else if ast.Tag == Prim {
		return "#<procedure: " + ast.Val.(Primitive).Name + ">"
	}
	return fmt.Sprintf("%v", ast.Val)
}

func (ast *AST) String() string {
	switch ast.Tag {
	case List, Quoted, Fn, Table, Prim:
		return ast.Pretty()
	case Char:
		return string(ast.Val.(rune))
	}
	return fmt.Sprintf("%v", ast.Val)
}

type Env struct {
	parent *Env
	Values map[string]*AST
}

type PrimFn func([]*AST) (*AST, error)

type Primitive struct {
	Name string
	Fn   PrimFn
}

func NewEnv(parent *Env) Env {
	return Env{parent, make(map[string]*AST)}
}

func ParentEnv() Env {
	return NewEnv(nil)
}

func (e *Env) Lookup(elem string) (*AST, error) {
	res, ok := e.Values[elem]

	if !ok {
		if e.parent == nil {
			return nil, fmt.Errorf("Symbol not found: %s", elem)
		}

		return e.parent.Lookup(elem)
	}

	return res, nil
}

type Func struct {
	Params []string
	Rest   *string
	Opt    *ordered_map.OrderedMap
	Body   *AST
	Env    Env
}

func (f *Func) HasRest() bool {
	return f.Rest != nil
}

func (f *Func) HasOpt() bool {
	return f.Opt.Len() != 0
}