Add const fn

xls/codegen/vast/translate_vast_to_dslx.cc

@@ -515,7 +515,7 @@ class VastToDslxTranslator {
         CreateNodeSpan(function), fn_name,
         /*parametric_bindings=*/std::vector<dslx::ParametricBinding*>(), params,
         fn_type, block, dslx::FunctionTag::kNormal,
-        /*is_public=*/true, /*is_test_utility=*/false);
+        /*is_public=*/true, /*is_const=*/false, /*is_test_utility=*/false);
     XLS_RETURN_IF_ERROR(module().AddTop(fn, /*make_collision_error=*/nullptr));
     fn_name->set_definer(fn);
     return fn;

xls/dslx/fmt/ast_fmt.cc

@@ -2173,6 +2173,12 @@ DocRef Formatter::Format(const Function& n, bool is_test) {
     signature_pieces.push_back(arena_.Make(Keyword::kPub));
     signature_pieces.push_back(arena_.space());
   }
+
+  if (n.is_const()) {
+    signature_pieces.push_back(arena_.Make(Keyword::kConst));
+    signature_pieces.push_back(arena_.space());
+  }
+
   signature_pieces.push_back(arena_.Make(Keyword::kFn));
   signature_pieces.push_back(arena_.space());
   signature_pieces.push_back(arena_.MakeText(n.identifier()));

xls/dslx/fmt/ast_fmt_test.cc

@@ -233,8 +233,9 @@ class FunctionFmtTest : public testing::Test {
       bindings_.Add(name, parser_->module().GetOrCreateBuiltinNameDef(name));
     }
 
-    XLS_ASSIGN_OR_RETURN(
-        f_, parser_->ParseFunction(Pos(), /*is_public=*/false, bindings_));
+    XLS_ASSIGN_OR_RETURN(f_,
+                         parser_->ParseFunction(Pos(), /*is_public=*/false,
+                                                /*is_const=*/false, bindings_));
     Comments comments = Comments::Create(scanner_->comments());
 
     DocRef doc = Formatter(comments, arena_).Format(*f_);

xls/dslx/frontend/ast.cc

@@ -2246,7 +2246,7 @@ Function::Function(Module* owner, Span span, NameDef* name_def,
                    std::vector<ParametricBinding*> parametric_bindings,
                    std::vector<Param*> params, TypeAnnotation* return_type,
                    StatementBlock* body, FunctionTag tag, bool is_public,
-                   bool is_stub)
+                   bool is_const, bool is_stub)
     : AstNode(owner),
       span_(std::move(span)),
       name_def_(name_def),
@@ -2256,6 +2256,7 @@ Function::Function(Module* owner, Span span, NameDef* name_def,
       body_(body),
       tag_(tag),
       is_public_(is_public),
+      is_const_(is_const),
       is_stub_(is_stub) {
   for (const ParametricBinding* pb : parametric_bindings_) {
     CHECK(parametric_keys_.insert(pb->identifier()).second)
@@ -2330,16 +2331,18 @@ std::string Function::ToString() const {
   }
   std::string pub_str = is_public() && !is_stub_ ? "pub " : "";
   std::string annotation_str;
+  std::string const_str = is_const() ? "const " : "";
 
   if (is_test_utility()) {
     annotation_str = absl::StrFormat("#[%s]\n", kCfgTestAttr);
   } else if (extern_verilog_module_.has_value()) {
     annotation_str = absl::StrFormat("#[extern_verilog(\"%s\")]\n",
                                      extern_verilog_module_->code_template());
   }
-  return absl::StrFormat("%s%sfn %s%s(%s)%s%s", annotation_str, pub_str,
-                         name_def_->ToString(), parametric_str, params_str,
-                         return_type_str, is_stub_ ? ";" : body_->ToString());
+  return absl::StrFormat("%s%s%sfn %s%s(%s)%s%s", annotation_str, pub_str,
+                         const_str, name_def_->ToString(), parametric_str,
+                         params_str, return_type_str,
+                         is_stub_ ? ";" : body_->ToString());
 }
 
 std::string Function::ToUndecoratedString(std::string_view identifier) const {

xls/dslx/frontend/ast.h

@@ -48,7 +48,7 @@
 
 // Higher-order macro for all the Expr node leaf types (non-abstract).
 #define XLS_DSLX_EXPR_NODE_EACH(X) \
-  /* keep-sorted start */       \
+  /* keep-sorted start */          \
   X(AllOnesMacro)                  \
   X(Array)                         \
   X(Attr)                          \
@@ -2467,7 +2467,8 @@ class Function : public AstNode {
   Function(Module* owner, Span span, NameDef* name_def,
            std::vector<ParametricBinding*> parametric_bindings,
            std::vector<Param*> params, TypeAnnotation* return_type,
-           StatementBlock* body, FunctionTag tag, bool is_public, bool is_stub);
+           StatementBlock* body, FunctionTag tag, bool is_public, bool is_const,
+           bool is_stub);
 
   ~Function() override;
   AstNodeKind kind() const override { return AstNodeKind::kFunction; }
@@ -2504,6 +2505,7 @@ class Function : public AstNode {
   bool IsParametric() const { return !parametric_bindings_.empty(); }
   bool is_public() const { return is_public_; }
   bool is_test_utility() const { return is_test_utility_; }
+  bool is_const() const { return is_const_; }
   void set_test_utility(bool value) { is_test_utility_ = value; }
   void set_compiler_derived(bool value) { is_compiler_derived_ = value; }
   bool IsMethod() const;
@@ -2581,6 +2583,7 @@ class Function : public AstNode {
   std::optional<Impl*> impl_;
 
   const bool is_public_;
+  const bool is_const_;
   bool is_test_utility_ = false;  // Set by the parser on applying attributes.
   bool is_compiler_derived_ = false;  // Set by type inference upon deriving.
   bool is_stub_;

xls/dslx/frontend/ast_cloner.cc

@@ -416,7 +416,8 @@ class AstCloner : public AstNodeVisitor {
         CastIfNotVerbatim<StatementBlock*>(old_to_new_.at(n->body())));
     auto new_function = module(n)->Make<Function>(
         n->span(), new_name_def, new_parametric_bindings, new_params,
-        new_return_type, new_body, n->tag(), n->is_public(), n->IsStub());
+        new_return_type, new_body, n->tag(), n->is_public(), n->is_const(),
+        n->IsStub());
     new_function->set_test_utility(n->is_test_utility());
     if (n->extern_verilog_module().has_value()) {
       new_function->set_extern_verilog_module(*n->extern_verilog_module());

xls/dslx/frontend/ast_test.cc

@@ -233,7 +233,7 @@ TEST(AstTest, GetFuncParam) {
 
   Function* f = m.Make<Function>(fake_span, func_name_def, parametric_bindings,
                                  params, u32_type_annotation, block,
-                                 FunctionTag::kNormal, false, false);
+                                 FunctionTag::kNormal, false, false, false);
 
   XLS_ASSERT_OK_AND_ASSIGN(Param * found_param, f->GetParamByName("p"));
   EXPECT_EQ(found_param, params[0]);

xls/dslx/frontend/function_specializer.cc

@@ -346,7 +346,8 @@ absl::StatusOr<Function*> InsertFunctionSpecialization(
       function_span, new_name_def,
       /*parametric_bindings=*/std::vector<ParametricBinding*>{}, new_params,
       new_return_type, new_body, source_function->tag(),
-      source_function->is_public(), source_function->is_test_utility());
+      source_function->is_public(), source_function->is_const(),
+      source_function->is_test_utility());
   new_name_def->set_definer(new_function);
 
   if (source_function->extern_verilog_module().has_value()) {

xls/dslx/frontend/parser.cc

@@ -244,10 +244,10 @@ absl::Status Parser::ParseErrorStatus(const Span& span,
 }
 
 absl::StatusOr<Function*> Parser::ParseFunction(
-    const Pos& start_pos, bool is_public, Bindings& bindings,
+    const Pos& start_pos, bool is_public, bool is_const, Bindings& bindings,
     absl::flat_hash_map<std::string, Function*>* name_to_fn) {
-  XLS_ASSIGN_OR_RETURN(Function * f,
-                       ParseFunctionInternal(start_pos, is_public, bindings));
+  XLS_ASSIGN_OR_RETURN(Function * f, ParseFunctionInternal(start_pos, is_public,
+                                                           is_const, bindings));
   if (name_to_fn == nullptr) {
     return f;
   }
@@ -319,8 +319,8 @@ absl::StatusOr<Lambda*> Parser::ParseLambda(Bindings& bindings) {
   NameDef* fn_name_def = module_->Make<NameDef>(sp, "lambda_fn", nullptr);
   Function* fn =
       module_->Make<Function>(sp, fn_name_def, parametrics, params, return_type,
-                              body, FunctionTag::kLambda,
-                              /*is_public=*/false, /*is_stub=*/false);
+                              body, FunctionTag::kLambda, /*is_public=*/false,
+                              /*is_const*/ false, /*is_stub=*/false);
   return module_->Make<Lambda>(sp, fn);
 }
 
@@ -466,9 +466,19 @@ absl::StatusOr<std::unique_ptr<Module>> Parser::ParseModule(
     if (!module_member_start_pos.has_value()) {
       module_member_start_pos = GetPos();
     }
+
     XLS_ASSIGN_OR_RETURN(is_public, TryDropKeyword(Keyword::kPub));
     XLS_ASSIGN_OR_RETURN(const Token* peek, PeekToken());
 
+    bool is_const_fn = false;
+    if (peek->IsKeyword(Keyword::kConst)) {
+      XLS_ASSIGN_OR_RETURN(const Token* postpeek, PeekToken(1));
+      if (postpeek->IsKeyword(Keyword::kFn)) {
+        XLS_ASSIGN_OR_RETURN(is_const_fn, TryDropKeyword(Keyword::kConst));
+        XLS_ASSIGN_OR_RETURN(peek, PeekToken());
+      }
+    }
+
     if (peek->IsIdentifier(kConstAssertIdentifier)) {
       XLS_RETURN_IF_ERROR(verify_no_attributes());
       XLS_RETURN_IF_ERROR(verify_not_public(*peek));
@@ -498,9 +508,9 @@ absl::StatusOr<std::unique_ptr<Module>> Parser::ParseModule(
 
     switch (peek->GetKeyword()) {
       case Keyword::kFn: {
-        XLS_ASSIGN_OR_RETURN(Function * fn,
-                             ParseFunction(*module_member_start_pos, is_public,
-                                           *bindings, &name_to_fn));
+        XLS_ASSIGN_OR_RETURN(
+            Function * fn, ParseFunction(*module_member_start_pos, is_public,
+                                         is_const_fn, *bindings, &name_to_fn));
         XLS_ASSIGN_OR_RETURN(ModuleMember fn_or_wrapper,
                              ApplyFunctionAttributes(fn, pending_attributes));
         XLS_RETURN_IF_ERROR(
@@ -611,6 +621,7 @@ absl::StatusOr<std::unique_ptr<Module>> Parser::ParseModule(
         return top_level_error();
     }
     is_public = false;
+    is_const_fn = false;
     pending_attributes.clear();
     module_member_start_pos = std::nullopt;
   }
@@ -2215,7 +2226,8 @@ absl::StatusOr<Import*> Parser::ParseImport(Bindings& bindings) {
 }
 
 absl::StatusOr<Function*> Parser::ParseFunctionInternal(
-    const Pos& start_pos, bool is_public, Bindings& outer_bindings) {
+    const Pos& start_pos, bool is_public, bool is_const,
+    Bindings& outer_bindings) {
   XLS_ASSIGN_OR_RETURN(Token fn_tok, PopKeywordOrError(Keyword::kFn));
 
   // Do not add the function name to bindings until after the signature is
@@ -2260,9 +2272,10 @@ absl::StatusOr<Function*> Parser::ParseFunctionInternal(
   } else {
     XLS_ASSIGN_OR_RETURN(body, ParseBlockExpression(bindings));
   }
-  Function* f = module_->Make<Function>(
-      Span(start_pos, GetPos()), name_def, std::move(parametric_bindings),
-      params, return_type, body, FunctionTag::kNormal, is_public, stub);
+  Function* f = module_->Make<Function>(Span(start_pos, GetPos()), name_def,
+                                        std::move(parametric_bindings), params,
+                                        return_type, body, FunctionTag::kNormal,
+                                        is_public, is_const, stub);
   name_def->set_definer(f);
   return f;
 }
@@ -3150,7 +3163,7 @@ absl::StatusOr<Function*> Parser::ParseProcConfig(
   Function* config = module_->Make<Function>(
       block->span(), name_def, std::move(parametric_bindings),
       std::move(config_params), return_type, block, FunctionTag::kProcConfig,
-      is_public, /*is_stub=*/false);
+      is_public, /*is_const=*/false, /*is_stub=*/false);
   name_def->set_definer(config);
 
   return config;
@@ -3212,7 +3225,7 @@ absl::StatusOr<Function*> Parser::ParseProcNext(
   Function* next = module_->Make<Function>(
       span, name_def, std::move(parametric_bindings),
       std::vector<Param*>({state_param}), return_type, body,
-      FunctionTag::kProcNext, is_public, /*is_stub=*/false);
+      FunctionTag::kProcNext, is_public, /*is_const=*/false, /*is_stub=*/false);
   name_def->set_definer(next);
 
   return next;
@@ -3240,7 +3253,7 @@ absl::StatusOr<Function*> Parser::ParseProcInit(
   Function* init = module_->Make<Function>(
       span, name_def, std::move(parametric_bindings), std::vector<Param*>(),
       /*return_type=*/nullptr, body, FunctionTag::kProcInit, is_public,
-      /*is_stub=*/false);
+      /*is_const=*/false, /*is_stub=*/false);
   name_def->set_definer(init);
   return init;
 }
@@ -4103,7 +4116,7 @@ absl::StatusOr<Impl*> Parser::ParseImpl(const Pos& start_pos, bool is_public,
       XLS_ASSIGN_OR_RETURN(
           Function * function,
           ParseFunctionInternal(member_start_pos, next_is_public,
-                                impl_bindings));
+                                /*is_const=*/false, impl_bindings));
       members.push_back(function);
     } else {
       return ParseErrorStatus(
@@ -4143,7 +4156,7 @@ absl::StatusOr<Trait*> Parser::ParseTrait(const Pos& start_pos, bool is_public,
       XLS_ASSIGN_OR_RETURN(
           Function * function,
           ParseFunctionInternal(member_start_pos, /*is_public=*/true,
-                                trait_bindings));
+                                /*is_const=*/false, trait_bindings));
       members.push_back(function);
     } else {
       return ParseErrorStatus(peek->span(),

xls/dslx/frontend/parser.h

@@ -121,7 +121,7 @@ class Parser : public TokenParser {
   FileTable& file_table() { return scanner().file_table(); }
 
   absl::StatusOr<Function*> ParseFunction(
-      const Pos& start_pos, bool is_public, Bindings& bindings,
+      const Pos& start_pos, bool is_public, bool is_const, Bindings& bindings,
       absl::flat_hash_map<std::string, Function*>* name_to_fn = nullptr);
 
   absl::StatusOr<Function*> ParseImplFunction(const Pos& start_pos,
@@ -611,7 +611,7 @@ class Parser : public TokenParser {
 
   // Parses a function out of the token stream.
   absl::StatusOr<Function*> ParseFunctionInternal(const Pos& start_pos,
-                                                  bool is_public,
+                                                  bool is_public, bool is_const,
                                                   Bindings& outer_bindings);
 
   absl::StatusOr<std::vector<Param*>> ParseParamsInternal(