From b73031b4c97b5eac315a36c09a2472bcd663aa1d Mon Sep 17 00:00:00 2001
From: Rashid Harvey <rosk.abed@gmail.com>
Date: Tue, 16 Apr 2024 18:37:54 +0300
Subject: [PATCH] Added the possibility to parse expressions like \f.(\x.x a)
 correctly

---
 docs/index.html | 12 +++++++++++-
 pylambda.py     | 12 +++++++++++-
 2 files changed, 22 insertions(+), 2 deletions(-)

diff --git a/docs/index.html b/docs/index.html
index f7d5a0a..8235290 100644
--- a/docs/index.html
+++ b/docs/index.html
@@ -92,12 +92,14 @@
     but also does alpha-reduction which the Haskell-version cannot provide
     """
 
+    import re
     from itertools import chain
     from string import ascii_letters, whitespace
     from dataclasses import dataclass
 
     Char = str
     letters = set(ascii_letters)
+    whitespace_pattern = re.compile(r"^\s+")
 
     class Expression: 
         def all_vars(self)->set[Char]:
@@ -264,7 +266,15 @@
                     raise ParseError("Found a special character in the head of a function")
                 rest = ""
                 if " " in body:
-                    body, rest = body.split(None, 1)
+                    count = 0
+                    for i,c in enumerate(body):
+                        if c == ")":
+                            count -= 1
+                        elif c == "(":
+                            count += 1
+                        elif whitespace_pattern.match(body[i:]) and not count:
+                            body, rest = body[:i], body[i:].lstrip()
+                            break
                 result.append(Func(head, parse_expr(body)))
             elif first == "(":
                 closing = find_closing_paren(rest)
diff --git a/pylambda.py b/pylambda.py
index bf5f501..e656258 100644
--- a/pylambda.py
+++ b/pylambda.py
@@ -5,12 +5,14 @@
 but also does alpha-reduction which the Haskell-version cannot provide
 """
 
+import re
 from itertools import chain
 from string import ascii_letters, whitespace
 from dataclasses import dataclass
 
 Char = str
 letters = set(ascii_letters)
+whitespace_pattern = re.compile(r"^\s+")
 
 class Expression: 
     def all_vars(self)->set[Char]:
@@ -177,7 +179,15 @@ def _parse_expr(rest: str) -> list[Expression]:
                 raise ParseError("Found a special character in the head of a function")
             rest = ""
             if " " in body:
-                body, rest = body.split(None, 1)
+                count = 0
+                for i,c in enumerate(body):
+                    if c == ")":
+                        count -= 1
+                    elif c == "(":
+                        count += 1
+                    elif whitespace_pattern.match(body[i:]) and not count:
+                        body, rest = body[:i], body[i:].lstrip()
+                        break
             result.append(Func(head, parse_expr(body)))
         elif first == "(":
             closing = find_closing_paren(rest)