From e84d03d383aa477c9ff05bcd016e36e84b71873b Mon Sep 17 00:00:00 2001
From: Chris Novakovic <chris@chrisn.me.uk>
Date: Mon, 9 Jun 2025 12:05:44 +0100
Subject: [PATCH] Add `//` operator

This is equivalent to Python's `//` operator: the mathematical division
of the first operand by the second with the `floor` function applied to
the result.
---
 docs/grammar.txt                                     |  2 +-
 src/parse/asp/grammar.go                             |  5 ++++-
 src/parse/asp/interpreter_test.go                    | 12 ++++++++++++
 src/parse/asp/lexer.go                               | 10 +++++++++-
 src/parse/asp/objects.go                             |  3 +++
 .../asp/test_data/interpreter/floor_divide.build     |  9 +++++++++
 6 files changed, 38 insertions(+), 3 deletions(-)
 create mode 100644 src/parse/asp/test_data/interpreter/floor_divide.build

diff --git a/docs/grammar.txt b/docs/grammar.txt
index 1c852e6b2e..2c22e28d8b 100644
--- a/docs/grammar.txt
+++ b/docs/grammar.txt
@@ -37,5 +37,5 @@ comprehension = "for" Ident { "," Ident } "in" expression
 slice = "[" [ expression ] [ ":" expression ] "]";
 lambda = "lambda" [ lambda_arg { "," lambda_arg } ] ":" expression;
 lambda_arg = Ident [ "=" expression ];
-operator = ("+" | "-" | "*" | "/" | "%" | "<" | ">" | "and" | "or" |
+operator = ("+" | "-" | "*" | "/" | "//" | "%" | "<" | ">" | "and" | "or" |
             "is" | "is" "not" | "in" | "not" "in" | "==" | "!=" | ">=" | "<=" | "|");
diff --git a/src/parse/asp/grammar.go b/src/parse/asp/grammar.go
index 6c303e5255..18084cabc0 100644
--- a/src/parse/asp/grammar.go
+++ b/src/parse/asp/grammar.go
@@ -283,6 +283,8 @@ const (
 	Multiply Operator = '×'
 	// Divide implements division, currently only between integers
 	Divide Operator = '÷'
+	// FloorDivide implements floor division, currently only between integers
+	FloorDivide Operator = '⑊'
 	// Modulo implements % (including string interpolation)
 	Modulo Operator = '%'
 	// Negate is the unary negation operator (not exactly the same as Subtract)
@@ -339,7 +341,7 @@ func (o Operator) Precedence() int {
 	switch o {
 	case Negate:
 		return 4
-	case Multiply, Divide, Modulo:
+	case Multiply, Divide, FloorDivide, Modulo:
 		return 3
 	case Add, Subtract:
 		return 2
@@ -366,6 +368,7 @@ var operators = map[string]Operator{
 	"-":      Subtract,
 	"*":      Multiply,
 	"/":      Divide,
+	"//":     FloorDivide,
 	"%":      Modulo,
 	"<":      LessThan,
 	">":      GreaterThan,
diff --git a/src/parse/asp/interpreter_test.go b/src/parse/asp/interpreter_test.go
index 1bd016f435..578e9b3039 100644
--- a/src/parse/asp/interpreter_test.go
+++ b/src/parse/asp/interpreter_test.go
@@ -444,6 +444,18 @@ func TestDivide(t *testing.T) {
 	assert.EqualValues(t, -2, s.Lookup("k"))
 }
 
+func TestFloorDivide(t *testing.T) {
+	s, err := parseFile("src/parse/asp/test_data/interpreter/floor_divide.build")
+	assert.NoError(t, err)
+	assert.EqualValues(t, 2, s.Lookup("a"))
+	assert.EqualValues(t, -2, s.Lookup("b"))
+	assert.EqualValues(t, 0, s.Lookup("c"))
+	assert.EqualValues(t, -1, s.Lookup("d"))
+	assert.EqualValues(t, 0, s.Lookup("e"))
+	assert.EqualValues(t, 1, s.Lookup("f"))
+	assert.EqualValues(t, 5, s.Lookup("g"))
+}
+
 func TestFStringOptimisation(t *testing.T) {
 	s, stmts, err := parseFileToStatements("src/parse/asp/test_data/interpreter/fstring_optimisation.build")
 	require.NoError(t, err)
diff --git a/src/parse/asp/lexer.go b/src/parse/asp/lexer.go
index f0b0964824..20fe5ec14d 100644
--- a/src/parse/asp/lexer.go
+++ b/src/parse/asp/lexer.go
@@ -252,7 +252,15 @@ func (l *lex) nextToken() Token {
 			return Token{Type: LexOperator, Value: string([]byte{next, l.bytes[l.pos-1]}), Pos: pos}
 		}
 		fallthrough
-	case ',', '.', '%', '*', '|', '&', ':', '/':
+	case ',', '.', '%', '*', '|', '&', ':':
+		return Token{Type: rune(next), Value: string(next), Pos: pos}
+	case '/':
+		// Look ahead one byte to see if this is a floor division.
+		if l.bytes[l.pos] == '/' {
+			l.pos++
+			l.col++
+			return Token{Type: LexOperator, Value: string([]byte{next, l.bytes[l.pos-1]}), Pos: pos}
+		}
 		return Token{Type: rune(next), Value: string(next), Pos: pos}
 	case '#':
 		// Comment character, consume to end of line.
diff --git a/src/parse/asp/objects.go b/src/parse/asp/objects.go
index 7195200ccf..e2723d9538 100644
--- a/src/parse/asp/objects.go
+++ b/src/parse/asp/objects.go
@@ -4,6 +4,7 @@ import (
 	"encoding/json"
 	"fmt"
 	"iter"
+	"math"
 	"slices"
 	"sort"
 	"strconv"
@@ -221,6 +222,8 @@ func (i pyInt) Operator(operator Operator, operand pyObject) pyObject {
 			return i * o
 		case Divide:
 			return i / o
+		case FloorDivide:
+			return newPyInt(int(math.Floor(float64(i) / float64(o))))
 		case LessThan:
 			return newPyBool(i < o)
 		case GreaterThan:
diff --git a/src/parse/asp/test_data/interpreter/floor_divide.build b/src/parse/asp/test_data/interpreter/floor_divide.build
new file mode 100644
index 0000000000..1804987daa
--- /dev/null
+++ b/src/parse/asp/test_data/interpreter/floor_divide.build
@@ -0,0 +1,9 @@
+a = 4 // 2
+b = 4 // -2
+c = 0 // 2
+d = -1 // 2
+e = -1 // -2
+
+# Precedence tests
+f = 8 / 4 // 2  # (8 / 4) // 2: / and // have the same precedence
+g = 4 + 3 // 2  # 4 + (3 // 2): // has greater precedence than +