Add support for arithmetic assignment statements

4 files changed, 332 insertions, 2 deletions

diff --git a/GRAMMAR.txt b/GRAMMAR.txt
index 68177fc..a2011e6 100644
--- a/GRAMMAR.txt
+++ b/GRAMMAR.txt
@@ -17,7 +17,12 @@ declaration -> constant | mutable
 constant -> "const" IDENTIFIER "=" expression
 mutable -> "mut" IDENTIFIER "=" expression
 
-assignment -> IDENTIFIER "=" expression
+assignment -> assign | addassign | subassign | mulassign | divassign
+assign -> IDENTIFIER "=" expression
+addassign -> IDENTIFIER "+=" expression
+subassign -> IDENTIFIER "-=" expression
+mulassign -> IDENTIFIER "*=" expression
+divassign -> IDENTIFIER "/=" expression
 
 expression -> equality
 equality -> comparison ( ( "==" | "!=" ) comparison )*
diff --git a/expression.go b/expression.go
index 9752730..27f4ddb 100644
--- a/expression.go
+++ b/expression.go
@@ -134,6 +134,11 @@ func (m *mutStmt) markDecl() {}
 
 func (m *mutStmt) line() int { return m.ln }
 
+type assignmentStmt interface {
+	statementExpr
+	markAssign()
+}
+
 type assignStmt struct {
 	name  string
 	value exprExpr
@@ -144,10 +149,66 @@ func (a *assignStmt) markExpr() {}
 
 func (a *assignStmt) markStmt() {}
 
-func (a *assignStmt) markDecl() {}
+func (a *assignStmt) markAssign() {}
 
 func (a *assignStmt) line() int { return a.ln }
 
+type addAssignStmt struct {
+	name  string
+	value exprExpr
+	ln    int
+}
+
+func (a *addAssignStmt) markExpr() {}
+
+func (a *addAssignStmt) markStmt() {}
+
+func (a *addAssignStmt) markAssign() {}
+
+func (a *addAssignStmt) line() int { return a.ln }
+
+type subAssignStmt struct {
+	name  string
+	value exprExpr
+	ln    int
+}
+
+func (s *subAssignStmt) markExpr() {}
+
+func (s *subAssignStmt) markStmt() {}
+
+func (s *subAssignStmt) markAssign() {}
+
+func (s *subAssignStmt) line() int { return s.ln }
+
+type mulAssignStmt struct {
+	name  string
+	value exprExpr
+	ln    int
+}
+
+func (m *mulAssignStmt) markExpr() {}
+
+func (m *mulAssignStmt) markStmt() {}
+
+func (m *mulAssignStmt) markAssign() {}
+
+func (m *mulAssignStmt) line() int { return m.ln }
+
+type divAssignStmt struct {
+	name  string
+	value exprExpr
+	ln    int
+}
+
+func (d *divAssignStmt) markExpr() {}
+
+func (d *divAssignStmt) markStmt() {}
+
+func (d *divAssignStmt) markAssign() {}
+
+func (d *divAssignStmt) line() int { return d.ln }
+
 type exprExpr interface {
 	expression
 	markExprExpr()
diff --git a/generate.go b/generate.go
index 5701cb1..51eb7c0 100644
--- a/generate.go
+++ b/generate.go
@@ -128,6 +128,14 @@ func generateStatement(stmt statementExpr, w io.Writer) error {
 		return generateMutStmt(s, false, w)
 	case *assignStmt:
 		return generateAssignStmt(s, w)
+	case *addAssignStmt:
+		return generateAddAssignStmt(s, w)
+	case *subAssignStmt:
+		return generateSubAssignStmt(s, w)
+	case *mulAssignStmt:
+		return generateMulAssignStmt(s, w)
+	case *divAssignStmt:
+		return generateDivAssignStmt(s, w)
 	}
 
 	return nil
@@ -195,6 +203,94 @@ func generateAssignStmt(stmt *assignStmt, w io.Writer) error {
 	return nil
 }
 
+func generateAddAssignStmt(stmt *addAssignStmt, w io.Writer) error {
+	ok, mutable := isDeclared(stmt.name, false)
+	if !ok {
+		return errUndeclared{name: stmt.name, line: stmt.line()}
+	}
+	if !mutable {
+		return errImmutable{name: stmt.name, line: stmt.line()}
+	}
+
+	value, err := stmt.value.generate(w)
+	if err != nil {
+		return err
+	}
+
+	in, out := allocReg(), allocReg()
+
+	fmt.Fprintf(w, "%s =w loadl %%%s\n", in, stmt.name)
+	fmt.Fprintf(w, "%s =w add %s, %s\n", out, in, value)
+	fmt.Fprintf(w, "storew %s, %%%s\n", out, stmt.name)
+	return nil
+}
+
+func generateSubAssignStmt(stmt *subAssignStmt, w io.Writer) error {
+	ok, mutable := isDeclared(stmt.name, false)
+	if !ok {
+		return errUndeclared{name: stmt.name, line: stmt.line()}
+	}
+	if !mutable {
+		return errImmutable{name: stmt.name, line: stmt.line()}
+	}
+
+	value, err := stmt.value.generate(w)
+	if err != nil {
+		return err
+	}
+
+	in, out := allocReg(), allocReg()
+
+	fmt.Fprintf(w, "%s =w loadl %%%s\n", in, stmt.name)
+	fmt.Fprintf(w, "%s =w sub %s, %s\n", out, in, value)
+	fmt.Fprintf(w, "storew %s, %%%s\n", out, stmt.name)
+	return nil
+}
+
+func generateMulAssignStmt(stmt *mulAssignStmt, w io.Writer) error {
+	ok, mutable := isDeclared(stmt.name, false)
+	if !ok {
+		return errUndeclared{name: stmt.name, line: stmt.line()}
+	}
+	if !mutable {
+		return errImmutable{name: stmt.name, line: stmt.line()}
+	}
+
+	value, err := stmt.value.generate(w)
+	if err != nil {
+		return err
+	}
+
+	in, out := allocReg(), allocReg()
+
+	fmt.Fprintf(w, "%s =w loadl %%%s\n", in, stmt.name)
+	fmt.Fprintf(w, "%s =w mul %s, %s\n", out, in, value)
+	fmt.Fprintf(w, "storew %s, %%%s\n", out, stmt.name)
+	return nil
+}
+
+func generateDivAssignStmt(stmt *divAssignStmt, w io.Writer) error {
+	ok, mutable := isDeclared(stmt.name, false)
+	if !ok {
+		return errUndeclared{name: stmt.name, line: stmt.line()}
+	}
+	if !mutable {
+		return errImmutable{name: stmt.name, line: stmt.line()}
+	}
+
+	value, err := stmt.value.generate(w)
+	if err != nil {
+		return err
+	}
+
+	in, out := allocReg(), allocReg()
+
+	fmt.Fprintf(w, "%s =w loadl %%%s\n", in, stmt.name)
+	fmt.Fprintf(w, "%s =w div %s, %s\n", out, in, value)
+	fmt.Fprintf(w, "storew %s, %%%s\n", out, stmt.name)
+	return nil
+}
+
 func (e *equalityExpr) generate(w io.Writer) (string, error) {
 	lhs, err := e.lhs.generate(w)
 	if err != nil {
diff --git a/parse.go b/parse.go
index c6b21f0..f15ac0a 100644
--- a/parse.go
+++ b/parse.go
@@ -292,6 +292,14 @@ func parseStatementContents(toks *tokens) (statementExpr, error) {
 		return decl, nil
 	}
 
+	assignment, err := parseAssignment(toks)
+	if err != nil {
+		return nil, err
+	}
+	if assignment != nil {
+		return assignment, nil
+	}
+
 	return nil, nil
 }
 
@@ -330,6 +338,10 @@ func parseDeclaration(toks *tokens) (declStmt, error) {
 		return mutable, nil
 	}
 
+	return nil, nil
+}
+
+func parseAssignment(toks *tokens) (assignmentStmt, error) {
 	assignment, err := parseAssignStmt(toks)
 	if err != nil {
 		return nil, err
@@ -338,6 +350,38 @@ func parseDeclaration(toks *tokens) (declStmt, error) {
 		return assignment, nil
 	}
 
+	addAssign, err := parseAddAssignStmt(toks)
+	if err != nil {
+		return nil, err
+	}
+	if addAssign != nil {
+		return addAssign, nil
+	}
+
+	subAssign, err := parseSubAssignStmt(toks)
+	if err != nil {
+		return nil, err
+	}
+	if subAssign != nil {
+		return subAssign, nil
+	}
+
+	mulAssign, err := parseMulAssignStmt(toks)
+	if err != nil {
+		return nil, err
+	}
+	if mulAssign != nil {
+		return mulAssign, nil
+	}
+
+	divAssign, err := parseDivAssignStmt(toks)
+	if err != nil {
+		return nil, err
+	}
+	if divAssign != nil {
+		return divAssign, nil
+	}
+
 	return nil, nil
 }
 
@@ -430,6 +474,130 @@ func parseAssignStmt(toks *tokens) (*assignStmt, error) {
 	}, nil
 }
 
+func parseAddAssignStmt(toks *tokens) (*addAssignStmt, error) {
+	nameTok, ok := toks.consumeToken()
+	if !ok {
+		return nil, unexpectedEOF
+	}
+	if nameTok.kind != identifier {
+		toks.unreadToken()
+		return nil, nil
+	}
+
+	ok, err := toks.match(plusEquals)
+	if err != nil {
+		return nil, err
+	}
+	if !ok {
+		toks.unreadTokens(2)
+		return nil, nil
+	}
+
+	expr, err := parseExpression(toks)
+	if err != nil {
+		return nil, err
+	}
+
+	return &addAssignStmt{
+		name:  nameTok.value,
+		value: expr,
+		ln:    nameTok.line,
+	}, nil
+}
+
+func parseSubAssignStmt(toks *tokens) (*subAssignStmt, error) {
+	nameTok, ok := toks.consumeToken()
+	if !ok {
+		return nil, unexpectedEOF
+	}
+	if nameTok.kind != identifier {
+		toks.unreadToken()
+		return nil, nil
+	}
+
+	ok, err := toks.match(minusEquals)
+	if err != nil {
+		return nil, err
+	}
+	if !ok {
+		toks.unreadTokens(2)
+		return nil, nil
+	}
+
+	expr, err := parseExpression(toks)
+	if err != nil {
+		return nil, err
+	}
+
+	return &subAssignStmt{
+		name:  nameTok.value,
+		value: expr,
+		ln:    nameTok.line,
+	}, nil
+}
+
+func parseMulAssignStmt(toks *tokens) (*mulAssignStmt, error) {
+	nameTok, ok := toks.consumeToken()
+	if !ok {
+		return nil, unexpectedEOF
+	}
+	if nameTok.kind != identifier {
+		toks.unreadToken()
+		return nil, nil
+	}
+
+	ok, err := toks.match(starEquals)
+	if err != nil {
+		return nil, err
+	}
+	if !ok {
+		toks.unreadTokens(2)
+		return nil, nil
+	}
+
+	expr, err := parseExpression(toks)
+	if err != nil {
+		return nil, err
+	}
+
+	return &mulAssignStmt{
+		name:  nameTok.value,
+		value: expr,
+		ln:    nameTok.line,
+	}, nil
+}
+
+func parseDivAssignStmt(toks *tokens) (*divAssignStmt, error) {
+	nameTok, ok := toks.consumeToken()
+	if !ok {
+		return nil, unexpectedEOF
+	}
+	if nameTok.kind != identifier {
+		toks.unreadToken()
+		return nil, nil
+	}
+
+	ok, err := toks.match(slashEquals)
+	if err != nil {
+		return nil, err
+	}
+	if !ok {
+		toks.unreadTokens(2)
+		return nil, nil
+	}
+
+	expr, err := parseExpression(toks)
+	if err != nil {
+		return nil, err
+	}
+
+	return &divAssignStmt{
+		name:  nameTok.value,
+		value: expr,
+		ln:    nameTok.line,
+	}, nil
+}
+
 func parseExpression(toks *tokens) (exprExpr, error) {
 	return parseEquality(toks)
 }