Support nested lambdas.

PiperOrigin-RevId: 893702803

Author: erinzmoore

xls/dslx/frontend/semantics_analysis.cc

@@ -257,6 +257,9 @@ class SideEffectExpressionFinder : public AstNodeVisitorWithDefault {
 class CollectNameRefs : public AstNodeVisitorWithDefault {
  public:
   absl::Status HandleNameRef(const NameRef* node) override {
+    if (node->IsBuiltin()) {
+      return DefaultHandler(node);
+    }
     if (node->GetDefiner() == nullptr ||
         (node->GetDefiner()->kind() != AstNodeKind::kFunction &&
          node->GetDefiner()->kind() != AstNodeKind::kImport)) {
@@ -352,6 +355,7 @@ class ReplaceLambdaWithInvocation : public AstNodeVisitorWithDefault {
   //     map(arr, lambda_capture{x: x}.call)
   //   }
   absl::Status HandleLambda(const Lambda* node) override {
+    XLS_RETURN_IF_ERROR(DefaultHandler(node));
     Module* module = node->owner();
     Function* original_fn = node->function();
     Span span = node->span();
@@ -471,6 +475,9 @@ class ReplaceLambdaWithInvocation : public AstNodeVisitorWithDefault {
         -> std::optional<AstNode*> {
       if (node->kind() == AstNodeKind::kNameRef) {
         const NameRef* name_ref = absl::down_cast<const NameRef*>(node);
+        if (name_ref->IsBuiltin()) {
+          return std::nullopt;
+        }
         const auto* name_def = std::get<const NameDef*>(name_ref->name_def());
         if (name_def != nullptr && seen.contains(name_def)) {
           NameRef* self_nr = node->owner()->Make<NameRef>(

xls/dslx/type_system_v2/typecheck_module_v2_test.cc

@@ -9679,6 +9679,31 @@ const_assert!(ARR[1] == u16:2);
               TypecheckSucceeds(HasNodeWithType("ARR", "uN[16][6]")));
 }
 
+TEST(TypecheckV2Test, NestedLambdas) {
+  EXPECT_THAT(
+      R"(
+fn main() -> u32[4][5] {
+  let z = zero!<u32[4][5]>();
+  map(enumerate(z), | tup | {
+    let i = tup.0;
+    let arr = tup.1;
+    map(enumerate(arr), | tup2 | {
+      let j = tup2.0;
+      i + j
+    })
+  })
+}
+
+const_assert!(main() == [[u32:0, 1, 2, 3],
+        [u32:1, 2, 3, 4],
+        [u32:2, 3, 4, 5],
+        [u32:3, 4, 5, 6],
+        [u32:4, 5, 6, 7]]);
+
+)",
+      TypecheckSucceeds(HasNodeWithType("main", "() -> uN[32][4][5]")));
+}
+
 TEST(TypecheckV2Test, LambdaWithImplicitReturn) {
   EXPECT_THAT(
       R"(