muesli/ast.go

// Abstract Syntax tree for the MUESLI interpreter
package muesli

import (
	"fmt"
	"strings"
)

type AstKind int

const (
	AstKindProgram = AstKind(iota)
	AstKindStatements
	AstKindStatement
	AstKindSet
	AstKindGet
	AstKindTarget
	AstKindCommand
	AstKindArguments
	AstKindArgument
	AstKindExpression
	AstKindBlock
	AstKindParenthesis
	AstKindList
	AstKindCapture
	AstKindWordValue
	AstKindWord
	AstKindType
	AstKindValue
	AstKindEnd
	AstKindError
)

func (astkind AstKind) String() string {
	switch astkind {
	case AstKindProgram:
		return "AstKindProgram"
	case AstKindStatements:
		return "AstKindStatements"
	case AstKindStatement:
		return "AstKindStatement"
	case AstKindSet:
		return "AstKindSet"
	case AstKindGet:
		return "AstKindGet"
	case AstKindTarget:
		return "AstKindTarget"
	case AstKindCommand:
		return "AstKindCommand"
	case AstKindArguments:
		return "AstKindArguments"
	case AstKindArgument:
		return "AstKindArgument"
	case AstKindExpression:
		return "AstKindExpression"
	case AstKindBlock:
		return "AstKindBlock"
	case AstKindParenthesis:
		return "AstKindParenthesis"
	case AstKindList:
		return "AstKindList"
	case AstKindCapture:
		return "AstKindCapture"
	case AstKindWordValue:
		return "AstKindWordValue"
	case AstKindWord:
		return "AstKindWord"
	case AstKindType:
		return "AstKindType"
	case AstKindValue:
		return "AstKindValue"
	case AstKindEnd:
		return "AstKindEnd"
	case AstKindError:
		return "AstKindError"
	default:
		return "Unknown AstKind"
	}
}

/* AST node kind */
type Ast struct {
	Parent *Ast
	Child  *Ast
	Before *Ast
	After  *Ast
	AstKind
	*Token
}

func NewAst(kind AstKind, parent *Ast, token *Token) *Ast {
	child := &Ast{parent, nil, nil, nil, kind, token}
	return child
}

func (ast *Ast) LastSibling() *Ast {
	res := ast
	for res != nil && res.After != nil {
		res = res.After
	}
	return res
}

func (ast *Ast) LastChild() *Ast {
	return ast.Child.LastSibling()
}

/* Detaches, I.E removes this node and all it's children from the parent tree. */
func (ast *Ast) Remove() *Ast {
	parent := ast.Parent
	before := ast.Before
	after := ast.After
	if before != nil {
		before.After = after
	}
	if after != nil {
		after.Before = before
	}
	if parent != nil {
		/* Special case if ast is the first child of it's parent. */
		if ast == parent.Child {
			parent.Child = after
		}
	}
	ast.Parent = nil
	return ast
}

func (ast *Ast) InsertSibling(sibling *Ast) *Ast {
	after := ast.After
	ast.After = sibling
	sibling.Before = ast
	sibling.After = after
	if after != nil {
		after.Before = sibling
	}
	sibling.Parent = ast.Parent
	return sibling
}

func (ast *Ast) AppendSibling(sibling *Ast) *Ast {
	return ast.LastSibling().InsertSibling(sibling)
}

func (ast *Ast) AppendChild(child *Ast) *Ast {
	child.Parent = ast
	if ast.Child == nil {
		ast.Child = child
	} else {
		ast.Child.AppendSibling(child)
	}
	return child
}

func (ast *Ast) NewSibling(kind AstKind, token *Token) *Ast {
	sibling := NewAst(kind, ast.Parent, token)
	return ast.AppendSibling(sibling)
}

func (ast *Ast) NewChild(kind AstKind, token *Token) *Ast {
	sibling := NewAst(kind, ast.Parent, token)
	return ast.AppendChild(sibling)
}

func (ast *Ast) Walk(walker func(node *Ast) *Ast) *Ast {
	if found := walker(ast); found != nil {
		return found
	}
	if ast.Child != nil {
		if found := ast.Child.Walk(walker); found != nil {
			return found
		}
	}
	if ast.After != nil {
		if found := ast.After.Walk(walker); found != nil {
			return found
		}
	}
	return nil
}

func (ast *Ast) String() string {
	if ast.Token == nil {
		return fmt.Sprintf("Ast %s nil", ast.AstKind.String())
	}
	return fmt.Sprintf("Ast %s %v", ast.AstKind.String(), ast.Token.String())
}

func (ast *Ast) Display() {
	ast.Walk(func(node *Ast) *Ast {
		depth := node.Depth()
		fmt.Printf(strings.Repeat("--", depth))
		if node != nil {
			fmt.Printf("Ast: %s\n", node.String())
		} else {
			fmt.Printf("Ast: nil node\n")
		}
		return nil
	})
}

func (ast *Ast) Depth() int {
	var depth int = 0
	parent := ast.Parent
	for parent != nil {
		depth++
		parent = parent.Parent
	}
	return depth
}

func (ast *Ast) CountChildren() int {
	var count int = 0
	child := ast.Child
	for child != nil {
		count++
		child = child.After
	}
	return count
}

func (ast *Ast) Errors() []*Ast {
	res := make([]*Ast, 0)
	if ast == nil {
		return res
	}
	ast.Walk(func(node *Ast) *Ast {
		if node != nil && node.AstKind == AstKindError {
			res = append(res, node)
		}
		return nil
	})
	return res
}