From 8d5858d3d2bb80bb73e0c5c6d5a8aec6fdcd08be Mon Sep 17 00:00:00 2001
From: Greg Shuflin <greg.shuflin@protonmail.com>
Date: Sat, 20 Nov 2021 23:56:53 -0800
Subject: [PATCH] Support custom binops

---
 schala-lang/src/parsing/combinator.rs | 29 ++++++++++++++++++++++++---
 schala-lang/src/parsing/peg_parser.rs |  4 +++-
 schala-lang/src/parsing/test.rs       | 17 ++++++++++++++++
 3 files changed, 46 insertions(+), 4 deletions(-)

diff --git a/schala-lang/src/parsing/combinator.rs b/schala-lang/src/parsing/combinator.rs
index 87eaebc..8fca4ce 100644
--- a/schala-lang/src/parsing/combinator.rs
+++ b/schala-lang/src/parsing/combinator.rs
@@ -284,9 +284,32 @@ fn func_decl(input: Span) -> ParseResult<Declaration> {
 
 //TODO handle operators
 fn func_signature(input: Span) -> ParseResult<Signature> {
-    map(tuple((kw("fn"), identifier, formal_params, opt(type_anno))), |(_, name, params, type_anno)| {
-        Signature { name: rc_string(name.fragment()), operator: false, params, type_anno }
-    })(input)
+    let normal_fn = context("ordinary-fn", tuple((identifier, formal_params, opt(type_anno))));
+    let operator_fn = context(
+        "operator-fn",
+        tuple((delimited(tok(char('(')), operator, tok(char(')'))), formal_params, opt(type_anno))),
+    );
+
+    context(
+        "func-signature",
+        preceded(
+            kw("fn"),
+            cut(alt((
+                map(normal_fn, |(name, params, type_anno)| Signature {
+                    name: rc_string(name.fragment()),
+                    operator: false,
+                    params,
+                    type_anno,
+                }),
+                map(operator_fn, |(op, params, type_anno)| Signature {
+                    name: rc_string(op.sigil()),
+                    operator: true,
+                    params,
+                    type_anno,
+                }),
+            ))),
+        ),
+    )(input)
 }
 
 fn formal_params(input: Span) -> ParseResult<Vec<FormalParam>> {
diff --git a/schala-lang/src/parsing/peg_parser.rs b/schala-lang/src/parsing/peg_parser.rs
index b1cc6e7..c54386b 100644
--- a/schala-lang/src/parsing/peg_parser.rs
+++ b/schala-lang/src/parsing/peg_parser.rs
@@ -112,10 +112,12 @@ peg::parser! {
         rule func_declaration(parser: &mut Parser) -> Declaration =
             _ sig:func_signature(parser) __ body:block(parser) { Declaration::FuncDecl(sig, body) }
 
-        //TODO handle operators
         rule func_signature(parser: &mut Parser) -> Signature =
             _ "fn" _ name:identifier() "(" _ params:formal_params(parser) _ ")" _ type_anno:type_anno()? { Signature {
                 name: rc_string(name), operator: false, params, type_anno
+            } } /
+            _ "fn" _ "(" op:operator() ")" _ "(" _ params:formal_params(parser) _ ")" _ type_anno:type_anno()? { Signature {
+                name: rc_string(op), operator: true, params, type_anno
             } }
 
         rule formal_params(parser: &mut Parser) -> Vec<FormalParam> =
diff --git a/schala-lang/src/parsing/test.rs b/schala-lang/src/parsing/test.rs
index f72f701..8947e46 100644
--- a/schala-lang/src/parsing/test.rs
+++ b/schala-lang/src/parsing/test.rs
@@ -856,6 +856,23 @@ fn functions() {
     );
 }
 
+#[test]
+fn custom_operator() {
+    let source = "fn (!!)(lhs,rhs)";
+    assert_ast!(
+        source,
+        vec![stmt(StatementKind::Declaration(Declaration::FuncSig(Signature {
+            name: rc("!!"),
+            operator: true,
+            params: vec![
+                FormalParam { name: rc("lhs"), default: None, anno: None },
+                FormalParam { name: rc("rhs"), default: None, anno: None },
+            ],
+            type_anno: None
+        })))]
+    );
+}
+
 #[test]
 fn max_function_params() {
     let mut buf = "fn longfunc(".to_string();