add as an operator and for-each loops

compiler/ast/operators.oc

@@ -35,6 +35,7 @@ enum Operator {
     Equals
     GreaterThan
     GreaterThanEquals
+    In
     Index
     LeftShift
     LessThan
@@ -78,10 +79,11 @@ def Operator::from_operator_overload_str(s: str): Operator => match s {
     "<<=" => LeftShiftEquals
     ">>=" => RightShiftEquals
     "%" => Modulus
+    "in" => In
     else => Error
 }
 
-def Operator::from_token(type: TokenType): Operator => match type {
+def Operator::from_token(tok: &Token): Operator => match tok.type {
     Ampersand => BitwiseAnd
     And => And
     Caret => BitwiseXor
@@ -103,7 +105,11 @@ def Operator::from_token(type: TokenType): Operator => match type {
     SlashEquals => DivideEquals
     Star => Multiply
     StarEquals => MultiplyEquals
-    else => std::panic(`Unhandled token type in Operator::from_token: {type.str()}`)
+    Identifier => match tok.text {
+        "in" => In
+        else => std::panic(`Unhandled identifier in Operator::from_token: {tok.text}`)
+    }
+    else => std::panic(`Unhandled token type in Operator::from_token: {tok.type.str()}`)
 }
 
 def Operator::num_overload_params(this): u32 => match this {
@@ -114,7 +120,7 @@ def Operator::num_overload_params(this): u32 => match this {
     RightShiftEquals | BitwiseXor | DivideEquals | Equals | GreaterThan |
     GreaterThanEquals | Index | LeftShift | LessThan | LessThanEquals |
     Minus | MinusEquals | Modulus | Multiply | MultiplyEquals | Or | Plus |
-    NotEquals | PlusEquals | RightShift => 2
+    NotEquals | PlusEquals | RightShift | In => 2
 
     IndexAssign => 3
 

compiler/parser.oc

@@ -83,6 +83,8 @@ def Parser::token_is(&this, type: TokenType): bool {
     return .token().type == type
 }
 
+def Parser::token_is_identifier(&this, name: str): bool => .token().is_identifier(name)
+
 def Parser::peek_token_is(&this, off: u32, type: TokenType): bool {
     if .curr + off >= .tokens.size return false
     let tok = .tokens.at(.curr + off)
@@ -813,7 +815,7 @@ def Parser::parse_term(&this, end_type: TokenType): &AST {
             .token_is(TokenType::Percent) {
         if .token_is(end_type) break
         let op_tok = .consume(.token().type)
-        let op = Operator::from_token(op_tok.type)
+        let op = Operator::from_token(op_tok)
         let rhs = .parse_cast(end_type)
         lhs = AST::new_binop(op, lhs, rhs, op_tok.span)
     }
@@ -826,7 +828,7 @@ def Parser::parse_additive(&this, end_type: TokenType): &AST {
     while .token_is(TokenType::Plus) or .token_is(TokenType::Minus) {
         if .token_is(end_type) break
         let op_tok = .consume(.token().type)
-        let op = Operator::from_token(op_tok.type)
+        let op = Operator::from_token(op_tok)
         let rhs = .parse_term(end_type)
         lhs = AST::new_binop(op, lhs, rhs, op_tok.span)
     }
@@ -855,7 +857,7 @@ def Parser::parse_bw_and(&this, end_type: TokenType): &AST {
     while .token_is(TokenType::Ampersand) {
         if .token_is(end_type) break
         let op_tok = .consume(.token().type)
-        let op = Operator::from_token(op_tok.type)
+        let op = Operator::from_token(op_tok)
         let rhs = .parse_shift(end_type)
         lhs = AST::new_binop(op, lhs, rhs, op_tok.span)
     }
@@ -867,7 +869,7 @@ def Parser::parse_bw_xor(&this, end_type: TokenType): &AST {
     while .token_is(TokenType::Caret) {
         if .token_is(end_type) break
         let op_tok = .consume(.token().type)
-        let op = Operator::from_token(op_tok.type)
+        let op = Operator::from_token(op_tok)
         let rhs = .parse_bw_and(end_type)
         lhs = AST::new_binop(op, lhs, rhs, op_tok.span)
     }
@@ -879,7 +881,7 @@ def Parser::parse_bw_or(&this, end_type: TokenType): &AST {
     while .token_is(TokenType::Line) {
         if .token_is(end_type) break
         let op_tok = .consume(.token().type)
-        let op = Operator::from_token(op_tok.type)
+        let op = Operator::from_token(op_tok)
         let rhs = .parse_bw_xor(end_type)
         lhs = AST::new_binop(op, lhs, rhs, op_tok.span)
     }
@@ -896,7 +898,9 @@ def Parser::parse_relational(&this, end_type: TokenType): &AST {
             .token_is(TokenType::LessThanEquals) or
             .token_is(TokenType::GreaterThanEquals) or
             .token_is(TokenType::EqualEquals) or
-            .token_is(TokenType::NotEquals) {
+            .token_is(TokenType::NotEquals) or
+            .token_is_identifier("in") {
+
         if .token_is(end_type) break
 
         let done = match .token().type {
@@ -919,7 +923,7 @@ def Parser::parse_relational(&this, end_type: TokenType): &AST {
         let tok = operators.at(i)
         let lhs = operands.at(i)
         let rhs = operands.at(i+1)
-        let op = AST::new_binop(Operator::from_token(tok.type), lhs, rhs, tok.span)
+        let op = AST::new_binop(Operator::from_token(tok), lhs, rhs, tok.span)
         if root? {
             root = AST::new_binop(And, root, op, tok.span)
         } else {
@@ -947,7 +951,7 @@ def Parser::parse_logical_and(&this, end_type: TokenType): &AST {
     while .token_is(TokenType::And) {
         if .token_is(end_type) break
         let op_tok = .consume(.token().type)
-        let op = Operator::from_token(op_tok.type)
+        let op = Operator::from_token(op_tok)
         let rhs = .parse_logical_not(end_type)
         lhs = AST::new_binop(op, lhs, rhs, op_tok.span)
     }
@@ -959,7 +963,7 @@ def Parser::parse_logical_or(&this, end_type: TokenType): &AST {
     while .token_is(TokenType::Or) {
         if .token_is(end_type) break
         let op_tok = .consume(.token().type)
-        let op = Operator::from_token(op_tok.type)
+        let op = Operator::from_token(op_tok)
         let rhs = .parse_logical_and(end_type)
         lhs = AST::new_binop(op, lhs, rhs, op_tok.span)
     }
@@ -990,7 +994,7 @@ def Parser::parse_expression(&this, end_type: TokenType): &AST {
             else => {
                 let tok = .consume(.token().type)
                 op_span = tok.span
-                yield Operator::from_token(tok.type)
+                yield Operator::from_token(tok)
             }
         }
 
@@ -1047,7 +1051,11 @@ def Parser::parse_if(&this): &AST {
 def Parser::parse_for_each(&this, start_span: Span): &AST {
 
     let name = .consume(Identifier)
-    .consume(TokenType::Colon)
+    if not (.token_is(TokenType::Colon) or .token_is_identifier("in")) {
+        .error(Error::new(.token().span, "Expected `:` of `in` after for-each loop variable"))
+        return AST::new(ASTType::Error, start_span)
+    }
+    .consume(.token().type)
 
     let expr = .parse_expression(end_type: TokenType::Newline)
     let iter_var_name = "__iter"
@@ -1137,7 +1145,10 @@ def Parser::parse_for_each(&this, start_span: Span): &AST {
 def Parser::parse_for(&this): &AST {
     let tok = .consume(TokenType::For)
 
-    if .token_is(Identifier) and .peek(1).type == TokenType::Colon {
+    if .token_is(Identifier) and (
+        .peek(1).type == TokenType::Colon or
+        .peek(1).is_identifier("in")
+    ) {
         return .parse_for_each(start_span: tok.span)
     }
 

compiler/passes/typechecker.oc

@@ -1198,6 +1198,14 @@ def TypeChecker::check_expression_helper(&this, node: &AST, hint: &Type): &Type
                     if not lhs? or not rhs? return null
                     return .check_binary_op(node, lhs.unaliased(), rhs.unaliased())
                 }
+                In => {
+                    // Swap the order of the operands for the potential overload
+                    let lhs = .check_expression(node.u.binary.lhs)
+                    let rhs = .check_expression(node.u.binary.rhs)
+                    if not lhs? or not rhs? return .get_base_type(Bool, node.span)
+                    .find_and_replace_overloaded_binary_op(In, node, node.u.binary.rhs, node.u.binary.lhs)
+                    return .get_base_type(Bool, node.span)
+                }
                 Index => return .check_index(node, hint, is_being_assigned: false)
                 Assignment => {
                     let lhs = .check_expression(node.u.binary.lhs)

compiler/tokens.oc

@@ -30,12 +30,16 @@ def Token::from_ident(text: str, span: Span): &Token {
 
 def Token::str(&this): str => `{.span.str()}: {.type.str()}`
 
-
 def Token::is_word(this): bool => match .type {
     TokenType::Identifier => true
     else => .type as u64 > TokenType::BEGIN_KEYWORDS as u64
 }
 
+def Token::is_identifier(this, name: str): bool => match .type {
+    TokenType::Identifier => name == .text
+    else => false
+}
+
 enum TokenType {
     // Other tokens
     AtSign

std/compact_map.oc

@@ -158,6 +158,7 @@ def Map::at(&this, key: K): V {
 
 def Map::size(&this): u32 => .items.size
 
+[operator "in"]
 def Map::contains(&this, key: K): bool {
     let hash = key.hash()
     let index = .get_index(key, hash)

std/map.oc

@@ -104,6 +104,7 @@ def Map::get(&this, key: K, defolt: V): V {
     return node.value
 }
 
+[operator "in"]
 //* Checks if the map contains the given key
 def Map::contains(&this, key: K): bool {
     return .get_item(key)?

std/set.oc

@@ -45,8 +45,10 @@ def Set::subtract(&this, other: &Set<T>) {
     .size = .map.size
 }
 
-def Set::iter(&this): Iterator<T> => Iterator<T>::make(this)
+[operator "in"]
 def Set::contains(&this, key: T): bool => .map.contains(key)
+
+def Set::iter(&this): Iterator<T> => Iterator<T>::make(this)
 def Set::is_empty(&this): bool => .map.size == 0
 
 def Set::free(&this) {

std/value.oc

@@ -150,6 +150,7 @@ def Value::push(&this, value: &Value) {
 [operator "+="] def Value::push_f32(&this, f: f32) => .push(Value::new_float(f as f64))
 [operator "+="] def Value::push_f64(&this, f: f64) => .push(Value::new_float(f))
 
+[operator "in"]
 def Value::contains(&this, key: str): bool {
     if .type != Dictionary return false
     return .u.as_dict.contains(key)

tests/in_operator.oc

@@ -0,0 +1,31 @@
+/// compile
+import std::set::{ Set }
+
+struct Foo {
+    x: u32
+    y: u32
+    z: u32
+}
+
+[operator "in"]
+def Foo::contains(this, x: u32): bool {
+    return .x == x or .y == x or .z == x
+}
+
+def main() {
+    let f = Foo(10, 20, 30)
+    println(`10 in f: {10 in f}`)
+    println(`15 in f: {15 in f}`)
+    println(`20 in f: {20 in f}`)
+
+    for c in "hello".chars() {
+        println(`{c}`)
+    }
+
+    let s = Set<u32>::new()
+    s += 10
+    let c = 10 in s
+
+    let j = 30
+    let d = j in s
+}

tests/set_operator_overloads.oc

@@ -11,7 +11,7 @@ def main() {
     s += 20
     s += 25
 
-    assert s.contains(5)
+    assert 5 in s
     assert s.size == 5
 
     let ns = Set<u32>::new()
@@ -23,21 +23,21 @@ def main() {
     s -= ns
 
     assert s.size == 2
-    assert s.contains(10)
-    assert s.contains(20)
-    assert not s.contains(5)
-    assert not s.contains(15)
-    assert not s.contains(25)
+    assert 10 in s
+    assert 20 in s
+    assert not 5 in s
+    assert not 15 in s
+    assert not 25 in s
 
     s += ns
     s -= 5
 
     assert s.size == 4
-    assert s.contains(10)
-    assert s.contains(15)
-    assert s.contains(20)
-    assert s.contains(25)
-    assert not s.contains(5)
+    assert 10 in s
+    assert 15 in s
+    assert 20 in s
+    assert 25 in s
+    assert not 5 in s
 
     println("Ok")
 }