Improve Python generation for do and letfn* forms (#795)

Fixes #794

Author: chrisrink10

CHANGELOG.md

@@ -18,7 +18,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
  * Allow vars to be callable to adhere to Clojure conventions (#767)
  * Adjust input path compatibility in `basilisp.core/load` input path to be relative to the namespace or the root path (#782)
  * No longer warn on unused bindings when their name begins with `_` (#756)
- * Improve the Python generation for `if` and `let*` forms to avoid unnecessary extra assignments (#793)
+ * Improve the Python generation for `do`, `if`, `let*`, and `letfn*` forms to avoid unnecessary extra assignments (#793, #794)
 
 ### Fixed
  * Fix issue with `(count nil)` throwing an exception (#759)

src/basilisp/lang/compiler/generator.py

@@ -1520,13 +1520,18 @@ def _do_to_py_ast(ctx: GeneratorContext, node: Do) -> GeneratedPyAST:
     )
 
     fn_body_ast: List[ast.AST] = []
-    do_result_name = genname(_DO_PREFIX)
     fn_body_ast.extend(map(statementize, body_ast.dependencies))
-    fn_body_ast.append(
-        ast.Assign(
-            targets=[ast.Name(id=do_result_name, ctx=ast.Store())], value=body_ast.node
+
+    if isinstance(body_ast.node, ast.Name):
+        do_result_name = body_ast.node.id
+    else:
+        do_result_name = genname(_DO_PREFIX)
+        fn_body_ast.append(
+            ast.Assign(
+                targets=[ast.Name(id=do_result_name, ctx=ast.Store())],
+                value=body_ast.node,
+            )
         )
-    )
 
     return GeneratedPyAST(
         node=ast.Name(id=do_result_name, ctx=ast.Load()), dependencies=fn_body_ast
@@ -2311,6 +2316,7 @@ def _let_to_py_ast(ctx: GeneratorContext, node: Let) -> GeneratedPyAST:
         if node.env.pos == NodeSyntacticPosition.EXPR:
             if isinstance(body_ast.node, ast.Name):
                 let_result_node = body_ast.node
+                assert isinstance(let_result_node.ctx, ast.Load)
             else:
                 let_result_name = genname(_LET_RESULT_PREFIX)
                 let_result_node = ast.Name(id=let_result_name, ctx=ast.Load())
@@ -2320,10 +2326,7 @@ def _let_to_py_ast(ctx: GeneratorContext, node: Let) -> GeneratedPyAST:
                         value=body_ast.node,
                     )
                 )
-            return GeneratedPyAST(
-                node=let_result_node,
-                dependencies=let_body_ast,
-            )
+            return GeneratedPyAST(node=let_result_node, dependencies=let_body_ast)
         else:
             let_body_ast.append(body_ast.node)
             return GeneratedPyAST(node=_noop_node(), dependencies=let_body_ast)
@@ -2356,21 +2359,23 @@ def _letfn_to_py_ast(ctx: GeneratorContext, node: LetFn) -> GeneratedPyAST:
                 )
             )
 
-        letfn_result_name = genname(_LETFN_RESULT_PREFIX)
         body_ast = _synthetic_do_to_py_ast(ctx, node.body)
         letfn_body_ast.extend(map(statementize, body_ast.dependencies))
 
         if node.env.pos == NodeSyntacticPosition.EXPR:
-            letfn_body_ast.append(
-                ast.Assign(
-                    targets=[ast.Name(id=letfn_result_name, ctx=ast.Store())],
-                    value=body_ast.node,
+            if isinstance(body_ast.node, ast.Name):
+                letfn_result_node = body_ast.node
+                assert isinstance(letfn_result_node.ctx, ast.Load)
+            else:
+                letfn_result_name = genname(_LETFN_RESULT_PREFIX)
+                letfn_result_node = ast.Name(id=letfn_result_name, ctx=ast.Load())
+                letfn_body_ast.append(
+                    ast.Assign(
+                        targets=[ast.Name(id=letfn_result_name, ctx=ast.Store())],
+                        value=body_ast.node,
+                    )
                 )
-            )
-            return GeneratedPyAST(
-                node=ast.Name(id=letfn_result_name, ctx=ast.Load()),
-                dependencies=letfn_body_ast,
-            )
+            return GeneratedPyAST(node=letfn_result_node, dependencies=letfn_body_ast)
         else:
             letfn_body_ast.append(body_ast.node)
             return GeneratedPyAST(node=_noop_node(), dependencies=letfn_body_ast)

tests/basilisp/compiler_test.py

@@ -2080,18 +2080,32 @@ def test_type_is_not_irecord(
                 lcompile(f"#{test_ns}.NewType{{:a 1 :b 2}}")
 
 
-def test_do(lcompile: CompileFn, ns: runtime.Namespace):
-    code = """
-    (do
-      (def first-name :Darth)
-      (def last-name "Vader"))
-    """
-    ns_name = ns.name
-    assert lcompile(code) == Var.find_in_ns(
-        sym.symbol(ns_name), sym.symbol("last-name")
-    )
-    assert lcompile("first-name") == kw.keyword("Darth")
-    assert lcompile("last-name") == "Vader"
+class TestDo:
+    def test_do(self, lcompile: CompileFn, ns: runtime.Namespace):
+        code = """
+        (do
+          (def first-name :Darth)
+          (def last-name "Vader"))
+        """
+        ns_name = ns.name
+        assert lcompile(code) == Var.find_in_ns(
+            sym.symbol(ns_name), sym.symbol("last-name")
+        )
+        assert lcompile("first-name") == kw.keyword("Darth")
+        assert lcompile("last-name") == "Vader"
+
+    def test_do_returning_name(self, lcompile: CompileFn, ns: runtime.Namespace):
+        assert (
+            lcompile(
+                """
+        (let* [a :a]
+          (do 
+            (def some-value a)
+            a))
+        """
+            )
+            == kw.keyword("a")
+        )
 
 
 class TestFunctionShadowName:
@@ -3453,6 +3467,10 @@ def test_letfn_may_have_empty_body(self, lcompile: CompileFn):
         assert None is lcompile("(letfn* [])")
         assert None is lcompile("(letfn* [a (fn* a [])])")
 
+    def test_let_return_bound_name(self, lcompile: CompileFn):
+        f = lcompile("(letfn* [a (fn* a [])] a)")
+        assert f() is None
+
     def test_letfn(self, lcompile: CompileFn):
         assert lcompile("(letfn* [a (fn* a [] 1)] (a))") == 1
         assert lcompile("(letfn* [a (fn* a [] 1) b (fn* b [] 2)] (b))") == 2