1 files changed, 103 insertions, 2 deletions
diff --git a/neb/parser.py b/neb/parser.py
index 22c51e4..360eac4 100644
--- a/neb/parser.py
+++ b/neb/parser.py
@@ -2,6 +2,99 @@ from .structs import *
 from .exceptions import ParseError
 
 
+def parseDef(token, prev, tokens):
+    if len(tokens) < 2:
+        raise ParseError("invalid def", token.line)
+    elif tokens[0].type_ != TokenType.SYMBOL:
+        raise ParseError("'def' must take a symbol", token.line)
+    sym, inc = parseSymbol(tokens[0], token, tokens[1:])
+    return NebDef(sym), inc + 1
+
+def parseLambda(token, prev, tokens):
+    # lambda [type] ( [arg] [&] ) [(expr)]
+    # func <name> [type] ( [arg] [&] ) [(expr)]
+    idx = 0
+    sym = None
+    return_type = Type(":any")
+    if token.type_ == TokenType.FUNC:
+        if tokens[idx].type_ != TokenType.SYMBOL:
+            raise ParseError("'func' must take a symbol", tokens[idx].line)
+        sym, counter = parseSymbol(tokens[idx], token, tokens[idx+1:])
+        idx += counter
+    else:
+        sym = Symbol("<lambda>", token.line)
+
+    if tokens[idx].type_ == TokenType.COLON:
+        return_type, counter = parseType(tokens[idx], token, tokens[idx+1:])
+        idx += counter
+
+    if tokens[idx].type_ != TokenType.OPEN_PAREN:
+        raise ParseError("expecting argument list", tokens[idx].line)
+
+    args, many, counter = parseFunctionArguments(tokens[idx:], tokens[idx-1], tokens[idx+1:])
+    idx += counter
+
+    body = []
+    while tokens[idx].type_ != TokenType.CLOSE_PAREN:
+        token = tokens[idx]
+        counter = 1
+        if token.type_ == TokenType.OPEN_PAREN:
+            expr, counter = parseExpression(token, prev, tokens[idx+1:])
+            body.append(expr)
+        elif token.type_ in (TokenType.FALSE, TokenType.TRUE, TokenType.STRING, TokenType.INT, TokenType.FLOAT):
+            lit, counter = parseLiteral(token, prev, tokens[idx+1:])
+            body.append(lit)
+        elif token.type_ == TokenType.COLON:
+            typ, counter = parseType(token, prev, tokens[idx+1:])
+            body.append(typ)
+        else:
+            sym, counter = parseSymbol(token, prev, tokens[idx+1:])
+            body.append(sym)
+        idx += counter
+        prev = token
+
+    return NebFuncDef(sym, return_type, args, many, body), idx + 1 
+
+def parseFunctionArguments(token, prev, tokens):
+    idx = 0
+    args = []
+    prev = token
+    many = None
+    prev_type = False
+    first = True
+    while tokens[idx].type_ != TokenType.CLOSE_PAREN:
+        token = tokens[idx]
+        if token.type_ == TokenType.SYMBOL:
+            if many is not None:
+                raise ParseError("& must be last argument", token.line)
+            sym, counter = parseSymbol(token, prev, tokens)
+            args.append(Arg(sym.name, TypeEnum.ANY))
+            prev_type = False
+            counter = 1
+        elif token.type_ == TokenType.MANY:
+            many = Arg("&", TypeEnum.ANY)
+            prev_type = False
+            counter = 1
+        elif token.type_ == TokenType.COLON and not prev_type and not first:
+            if prev_type:
+                raise ParseError("can't have two types in a row", token.line) 
+            if first:
+                raise ParseError("type can't be first", token.line) 
+            typ, counter = parseType(token, prev, tokens[idx+1:])
+            if many is None:
+                args[-1].type_ = typ
+            else:
+                many.type_ = typ
+            prev_type = True
+        else:
+            raise ParseError("invalid function signature", token)
+        first = False
+        idx += counter
+
+    # this should return function
+    return args, many, idx + 2  # parens
+
+
 def parseExpression(token, prev, tokens):
     idx = 0
     args = []
@@ -20,6 +113,12 @@ def parseExpression(token, prev, tokens):
         elif token.type_ == TokenType.COLON:
             expr, inc = parseType(token, prev, tokens[idx+1:])
             args.append(expr)
+        elif token.type_ == TokenType.DEF:
+            expr, inc = parseDef(token, prev, tokens[idx+1:])
+            args.append(expr)
+        elif token.type_ in (TokenType.LAMBDA, TokenType.FUNC):
+            expr, inc = parseLambda(token, prev, tokens[idx+1:])
+            args.append(expr)
         else:
             expr, inc = parseSymbol(token, prev, tokens[idx+1:])
             args.append(expr)
@@ -29,8 +128,10 @@ def parseExpression(token, prev, tokens):
     return Expr(args), idx + 2  # parens
 
 def parseSymbol(token, prev, tokens):
-    if token.text.startswith("'"):
-        return Symbol(token.text[1:], token.line, True), 1
+    if token.type_ == TokenType.APOSTROPHE:
+        if len(tokens) == 0 or tokens[0].type_ != TokenType.SYMBOL:
+            raise ParseError("quote must be followed by a symbol", token.line)
+        return Symbol(tokens[0].text, tokens[0].line, True), 2
     else:
         return Symbol(token.text, token.line), 1