typing.wren

// an excruciatingly overcomplicated calculator with memory written in Wren
// (made for learning how to do typechecking)
// run with the Wren CLI
// https://wren.io/cli

////////////////////////////////////////////////////////////////////

class Token {
	construct new(type, lexeme, line) {
		_type = type
		_lexeme = lexeme
		_line = line
	}

	type   { _type   }
	lexeme { _lexeme }
	line   { _line   }

	toString {
		var s = "%(_type) "
		if (_lexeme.count != 0) s = s + "'%(_lexeme)' "
		return s + "(line %(_line))"
	}
}

class Lexer {
	construct new(str) {
		if (!(str is String)) Fiber.abort("expected string as input")

		_input = str
		_start = _current = 0
		_line = 1

		if (!__keywords) __keywords = {
			"bool":  "BOOL",
			"num":   "NUM",
			"true":  "TRUE",
			"false": "FALSE",
		}

		if (!__chars) __chars = {
			"(": "LPAREN",
			")": "RPAREN",
			"+": "PLUS",
			"-": "MINUS",
			"*": "STAR",
			"/": "SLASH",
			"_": "UNDERSCORE",
			"!": "BANG",
			"=": "EQUALS",
			"^": "CARET",
			"<": "LESS",
			">": "GREATER",
			"@": "AT",
		}
	}

	isDigit_(c) {
		var b = c.bytes[0]
		return b >= 48 && b <= 57
	}

	isAlpha_(c) {
		var b = c.bytes[0]
		return b >= 97 && b <= 122
	}

	advance_() {
		var c = _input[_current]
		_current = _current + 1
		if (c == "\n") _line = _line + 1
		return c
	}

	peek_ {
		if (_current >= _input.count) return "\0"
		return _input[_current]
	}

	isAtEnd_ { _current >= _input.count }

	token_(type, text) { Token.new(type, text, _line) }
	token_(type) { token_(type, _input[_start..._current]) }

	number_() {
		while (isDigit_(peek_)) advance_()
		return token_("NUMBER")
	}

	word_() {
		while (isAlpha_(peek_) || isDigit_(peek_)) {
			advance_()
		}

		var id = _input[_start..._current]
		var kw = __keywords[id]
		if (kw) return token_(kw, "")
		return token_("IDENTIFIER", id)
	}

	next() {
		while (!isAtEnd_ && peek_.trim().count == 0) advance_()
		if (isAtEnd_) return token_("EOF", "")
		_start = _current

		var c = advance_()
		if (isDigit_(c)) return number_()
		if (isAlpha_(c)) return word_()

		if (__chars[c]) return token_(__chars[c])
		return token_("error", "unexpected character %(c)")
	}

	iterate(unused) {
		var tok = next()
		return tok.type != "EOF" ? tok : false
	}

	iteratorValue(tok) { tok }
}

////////////////////////////////////////////////////////////////////

class Node {
	visit_(v) { Fiber.abort("unimplemented visit_ for %(this.type)") }
}

class BinOp is Node {
	construct new(a, op, b) {
		_left = a
		_operator = op
		_right = b
	}
	left     { _left     }
	operator { _operator }
	right    { _right    }
	toString { "BinOp(%(_left), %(_operator), %(_right))" }
	visit_(v) { v.visitBinOp_(this) }
}

class UnOp is Node {
	construct new(op, b) {
		_operator = op
		_value = b
	}
	operator { _operator }
	value    { _value    }
	toString { "UnOp(%(_operator), %(_value))" }
	visit_(v) { v.visitUnOp_(this) }
}

class Statements is Node {
	construct new(stmts) {
		_statements = stmts
	}
	statements { _statements }
	toString { "Statements(%(_statements))" }
	visit_(v) { v.visitStatements_(this) }
}

class Assign is Node {
	construct new(type, id, exp) {
		_valueType = type
		_identifier = id
		_expression = exp
	}
	valueType  { _valueType  }
	identifier { _identifier }
	expression { _expression }
	toString { "Assign(%(_valueType), %(_identifier), %(_expression))"}
	visit_(v) { v.visitAssign_(this) }
}

class Variable is Node {
	construct new(id) {
		_identifier = id
	}
	identifier { _identifier }
	toString { "Variable(%(_identifier))" }
	visit_(v) { v.visitVariable_(this) }
}

class Literal is Node {
	construct new(value) { _value = value }
	value { _value }
	toString { "Literal(%(_value))" }
	visit_(v) { v.visitLiteral_(this) }
}

class Meta {}

class Type is Meta {
	construct new(id) {
		_identifier = id
	}

	identifier { _identifier }

	toString { "Type(%(_identifier))" }

	==(o) {
		if (!(o is Type)) Fiber.abort("attempt to compare Type to %(o.type)")
		return _identifier == o.identifier
	}

	!=(o) { !(this == o) }
}

var BOOLEAN_TYPE = Type.new("BOOLEAN")
var NUMBER_TYPE = Type.new("NUMBER")

////////////////////////////////////////////////////////////////////

class ParserError {
	construct new(token, message) {
		_token = token
		_message = message
	}

	token   { _token   }
	message { _message }

	toString { "error at %(_token.line): %(_message)" }
}

class Parser {
	construct new(l) {
		if (!((_l = l) is Lexer)) Fiber.abort("expected a lexer")
		_cur = _next = null
		_errors = []
	}

	error_(t, msg) {
		_errors.add(ParserError.new(t, msg))
		return false
	}

	errors { _errors }

	advance_() {
		_cur = _next
		var error = false
		while (true) {
			_next = _l.next()
			if (_next.type != "error") return !error
			error_(_next, _next.lexeme)
			error = true
		}
	}

	check_(types) {
		if (!(types is List)) return _cur.type == types
		return types.any {|t| _cur.type == t}
	}

	isAtEnd_ { check_("EOF") }

	consume_(t, msg) {
		var cur = _cur
		if (match_(t)) return cur
		return error_(cur, msg)
	}

	consume_(t) { consume_(t, "expected %(t) but got %(_cur.type)") }

	match_(t) {
		if (check_(t)) {
			advance_()
			return true
		}
		return false
	}

	// primary = IDENTIFIER / NUMBER / TRUE / FALSE / (LPAREN expression RPAREN)
	primary_() {
		var cls
		if (check_("IDENTIFIER")) {
			cls = Variable
		} else if (check_(["NUMBER", "TRUE", "FALSE"])) {
			cls = Literal
		} else if (match_("LPAREN")) {
			var n = expr_()
			if (!consume_("RPAREN", "expected ')' to close expression")) return
			return n
		} else {
			return error_(_cur, "expected expression, got %(_cur.type)")
		}
		var n = cls.new(_cur)
		advance_()
		return n
	}

	// prefix = *1(BANG / AT) primary
	prefix_() {
		if (check_(["BANG", "AT"])) {
			var t = _cur
			advance_()
			return UnOp.new(t, prefix_())
		}
		return primary_()
	}

	// I can't be bothered to make a Pratt parser

	// factor = 1*((star / slash) factor)
	factor_() {
		var left = prefix_()
		while (check_(["STAR", "SLASH"])) {
			var op = _cur
			advance_()
			var right = primary_()
			if (!right) return
			left = BinOp.new(left, op, right)
		}
		return left
	}

	// term = 1*((PLUS / MINUS) term)
	term_() {
		var left = factor_()
		while (check_(["PLUS", "MINUS"])) {
			var op = _cur
			advance_()
			var right = factor_()
			if (!right) return
			left = BinOp.new(left, op, right)
		}
		return left
	}

	// comparison = 1*((LESS / GREATER) comparison)
	comparison_() {
		var left = term_()
		while (check_(["LESS", "GREATER"])) {
			var op = _cur
			advance_()
			var right = term_()
			if (!right) return
			left = BinOp.new(left, op, right)
		}
		return left
	}

	// equality = 1*((EQUALS / CARET) equality)
	equality_() {
		var left = comparison_()
		while (check_(["EQUALS", "CARET"])) {
			var op = _cur
			advance_()
			var right = comparison_()
			if (!right) return
			left = BinOp.new(left, op, right)
		}
		return left
	}

	// end of code repetition

	// expression = equality
	expr_() { equality_() }

	// type = BOOL / NUM
	type_() {
		var t
		if (check_("BOOL")) {
			t = BOOLEAN_TYPE
		} else if (check_("NUM")) {
			t = NUMBER_TYPE
		} else {
			advance_()
			return error_(_cur, "expected type but got %(_cur.type)")
		}
		advance_()
		return t
	}

	// assign = type (IDENTIFIER/UNDERSCORE) expression
	assign_() {
		var t
		if (!(t = type_())) return
		var id
		if (!match_("UNDERSCORE")) {
			if (!(id = consume_("IDENTIFIER"))) return
		}
		var e
		if (!(e = expr_())) return
		return Assign.new(t, id, e)
	}

	// program = *assign eof
	program_() {
		var statements = []
		while (!isAtEnd_) {
			statements.add(assign_())
		}
		return Statements.new(statements)
	}

	parse() {
		if (!advance_()) return
		if (!advance_()) return

		var stmts = program_()
		consume_("EOF")
		return stmts
	}
}

////////////////////////////////////////////////////////////////////

class Visitor {
	visit_(n) { n.visit_(this) }

	visitBinOp_(n)    { Fiber.abort("no visit_ method for BinOp")    }
	visitUnOp_(n)     { Fiber.abort("no visit_ method for UnOp")     }
	visitAssign_(n)   { Fiber.abort("no visit_ method for Assign")   }
	visitVariable_(n) { Fiber.abort("no visit_ method for Variable") }
	visitLiteral_(n)  { Fiber.abort("no visit_ method for Literal")  }
}

////////////////////////////////////////////////////////////////////

class SemaError {
	construct new(node, message) {
		_node = node
		_message = message
	}

	node    { _node   }
	message { _message }

	toString { "error: %(_message)" } // TODO: use _node in error report somehow?
}

class Sema is Visitor {
	construct new() {
		_variables = {}
		_errors = []
	}

	error_(n, msg) {
		_errors.add(SemaError.new(n, msg))
		return
	}

	errors { _errors }

	check(ast) {
		_errors.clear()
		return visit_(ast)
	}

	visitStatements_(n) {
		for (s in n.statements) if (!visit_(s)) return
		return true
	}

	visitAssign_(n) {
		var type = n.valueType

		if (n.identifier) {
			var id = n.identifier.lexeme
			if (_variables[id] && _variables[id] != type) {
				return error_(n, "attempt to reassign to '%(id)' with a declaration of a different type")
			}
			_variables[id] = type
		}

		var t
		if (!(t = visit_(n.expression))) return
		if (type != t) {
			var s = n.identifier ?
				"type mismatch in assignment to '%(n.identifier.lexeme)'" :
				"type assertion failed"

			return error_(n, "%(s): expected %(type.identifier), got %(t.identifier)")
		}

		return type
	}

	compatibleOp_(op, v) { compatibleOp_(v, op, null) }

	compatibleOp_(a, op, b) {
		// unary operations
		if (op == "AT"  ) return a == NUMBER_TYPE ? BOOLEAN_TYPE : NUMBER_TYPE
		if (op == "BANG") return BOOLEAN_TYPE

		if (!b) Fiber.abort("exhausted all unary op possibilities for '%(op)', none matched")

		// optimization: if the types are different, it must be an equals comparison
		if (a != b && op == "EQUALS") return BOOLEAN_TYPE

		// binary operations
		if (["PLUS", "MINUS", "STAR", "SLASH", "LESS", "GREATER"].contains(op) &&
		    a == NUMBER_TYPE) return NUMBER_TYPE

		// equality is compatible with any types
		if (["EQUALS", "CARET"].contains(op)) return BOOLEAN_TYPE

		return null // the operators matched no type
	}

	visitUnOp_(n) {
		var val = visit_(n.value)
		var t = compatibleOp_(n.operator.type, val)
		if (!t) return error_(n, "type mismatch in unary '%(n.operator.lexeme)'")
		return t
	}

	visitBinOp_(n) {
		var a = visit_(n.left)
		var b = visit_(n.right)
		var t = compatibleOp_(a, n.operator.type, b)
		if (!t) return error_(n, "type mismatch in binary '%(n.operator.lexeme)'")
		return t
	}

	visitVariable_(n) {
		var t = _variables[n.identifier.lexeme]
		if (!t) return error_(n, "attempt to use undefined variable '%(n.identifier)'")
		return t
	}

	visitLiteral_(n) {
		if (n.value.type == "NUMBER") return NUMBER_TYPE
		if (["TRUE", "FALSE"].contains(n.value.type)) return BOOLEAN_TYPE
		Fiber.abort("unreachable")
	}
}

////////////////////////////////////////////////////////////////////

// ...and, as an extra, an interpreter

class Value {
	toString { value.toString }
}

// itz monad time

class Number is Value {
	construct new(n) { _value = n }
	value { _value }

	!=(o) { Number.new(_value != o.value) }
	==(o) { Number.new(_value == o.value) }
	+(o)  { Number.new(_value +  o.value) }
	-(o)  { Number.new(_value -  o.value) }
	*(o)  { Number.new(_value *  o.value) }
	/(o)  { Number.new(_value /  o.value) }
	>(o)  { Number.new(_value >  o.value) }
	<(o)  { Number.new(_value <  o.value) }
	!     { Boolean.new(false)            }

	applyAt { _value != 0 }
}

class Boolean is Value {
	construct new(n) { _value = n }
	value { _value }

	==(o) { Boolean.new(_value == o.value) }
	!=(o) { Boolean.new(_value != o.value) }
	!     { Boolean.new(!_value)           }

	applyAt { _value ? 1 : 0 }
}

// end of monad time

class Interpreter is Visitor {
	construct new() {
		_variables = {}
		_errors = [] // currently useless field is currently useless
	}

	errors { _errors }

	interpret(ast) {
		_errors.clear()
		return visit_(ast)
	}

	visitStatements_(n) {
		var final
		for (s in n.statements) final = visit_(s)
		return final
	}

	visitUnOp_(n) {
		var value = visit_(n.value)
		var op = n.operator.type

		if (op == "BANG") return !value
		if (op == "AT"  ) return value.applyAt

		Fiber.abort("unreachable")
	}

	visitBinOp_(n) {
		var left = visit_(n.left)
		var right = visit_(n.right)
		var op = n.operator.type

		if (op == "PLUS"   ) return left + right
		if (op == "MINUS"  ) return left - right
		if (op == "STAR"   ) return left * right
		if (op == "SLASH"  ) return left / right
		if (op == "LESS"   ) return left < right
		if (op == "GREATER") return left > right
		if (op == "CARET"  ) return left != right
		if (op == "EQUALS" ) return left == right

		Fiber.abort("unreachable")
	}

	visitLiteral_(n) {
		if (n is Boolean) return Boolean.new(n.value.type == "TRUE")
		return Number.new(Num.fromString(n.value.lexeme))
	}

	visitAssign_(n)   {
		var v = visit_(n.expression)
		// if not underscore
		if (n.identifier) _variables[n.identifier.lexeme] = v
		return v
	}

	visitVariable_(n) { _variables[n.identifier.lexeme] }
}

////////////////////////////////////////////////////////////////////

import "io" for Stdin, Stdout

System.print("send 'quit' to... you know...")
var i = Interpreter.new()
var s = Sema.new()
while (true) {
	System.write("> ")
	Stdout.flush()
	var input = Stdin.readLine()
	if (input.trim() == "quit") break

	var p = Parser.new(Lexer.new(input))
	var n = p.parse()
	if (p.errors.count > 0) {
		p.errors.each{|e| System.print(e)}
		continue
	}

	if (!s.check(n)) {
		s.errors.each{|e| System.print(e)}
		continue
	}

	var v = i.interpret(n)
	if (i.errors.count > 0) {
		i.errors.each{|e| System.print(e)}
		continue
	}

	System.print(v)
}