Skip to content

Commit fe5d2c5

Browse files
asteriteasterite
committed
feat: add mutating FunctionDefinition functions
1 parent 5f9996e commit fe5d2c5

9 files changed

Lines changed: 248 additions & 9 deletions

File tree

‎compiler/noirc_frontend/src/elaborator/patterns.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -39,7 +39,7 @@ impl<'context> Elaborator<'context> {
3939

4040
/// Equivalent to `elaborate_pattern`, this version just also
4141
/// adds any new DefinitionIds that were created to the given Vec.
42-
pub(super) fn elaborate_pattern_and_store_ids(
42+
pub fn elaborate_pattern_and_store_ids(
4343
&mut self,
4444
pattern: Pattern,
4545
expected_type: Type,

‎compiler/noirc_frontend/src/hir/comptime/errors.rs‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -185,6 +185,9 @@ pub enum InterpreterError {
185185
FailedToResolveTraitDefinition {
186186
location: Location,
187187
},
188+
CannotMutateFunction {
189+
location: Location,
190+
},
188191

189192
Unimplemented {
190193
item: String,
@@ -255,6 +258,7 @@ impl InterpreterError {
255258
| InterpreterError::TraitDefinitionMustBeAPath { location }
256259
| InterpreterError::FailedToResolveTraitDefinition { location }
257260
| InterpreterError::FailedToResolveTraitBound { location, .. } => *location,
261+
InterpreterError::CannotMutateFunction { location, .. } => *location,
258262

259263
InterpreterError::FailedToParseMacro { error, file, .. } => {
260264
Location::new(error.span(), *file)
@@ -516,6 +520,10 @@ impl<'a> From<&'a InterpreterError> for CustomDiagnostic {
516520
let msg = "Failed to resolve to a trait definition".to_string();
517521
CustomDiagnostic::simple_error(msg, String::new(), location.span)
518522
}
523+
InterpreterError::CannotMutateFunction { location } => {
524+
let msg = "Cannot mutate function at this point".to_string();
525+
CustomDiagnostic::simple_error(msg, String::new(), location.span)
526+
}
519527
}
520528
}
521529
}

‎compiler/noirc_frontend/src/hir/comptime/interpreter/builtin.rs‎

Lines changed: 183 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -7,19 +7,24 @@ use acvm::{AcirField, FieldElement};
77
use builtin_helpers::{
88
check_argument_count, check_one_argument, check_three_arguments, check_two_arguments,
99
get_function_def, get_module, get_quoted, get_slice, get_trait_constraint, get_trait_def,
10-
get_type, get_u32, hir_pattern_to_tokens,
10+
get_tuple, get_type, get_u32, hir_pattern_to_tokens,
1111
};
1212
use chumsky::Parser;
1313
use iter_extended::{try_vecmap, vecmap};
1414
use noirc_errors::Location;
1515
use rustc_hash::FxHashMap as HashMap;
1616

1717
use crate::{
18-
ast::IntegerBitSize,
18+
ast::{
19+
FunctionKind, FunctionReturnType, IntegerBitSize, UnresolvedType, UnresolvedTypeData,
20+
Visibility,
21+
},
1922
hir::comptime::{errors::IResult, value::add_token_spans, InterpreterError, Value},
23+
hir_def::function::FunctionBody,
2024
macros_api::{ModuleDefId, NodeInterner, Signedness},
25+
node_interner::DefinitionKind,
2126
parser,
22-
token::Token,
27+
token::{SpannedToken, Token},
2328
QuotedType, Shared, Type,
2429
};
2530

@@ -43,6 +48,11 @@ impl<'local, 'context> Interpreter<'local, 'context> {
4348
"function_def_name" => function_def_name(interner, arguments, location),
4449
"function_def_parameters" => function_def_parameters(interner, arguments, location),
4550
"function_def_return_type" => function_def_return_type(interner, arguments, location),
51+
"function_def_set_body" => function_def_set_body(self, arguments, location),
52+
"function_def_set_parameters" => function_def_set_parameters(self, arguments, location),
53+
"function_def_set_return_type" => {
54+
function_def_set_return_type(self, arguments, location)
55+
}
4656
"module_functions" => module_functions(self, arguments, location),
4757
"module_is_contract" => module_is_contract(self, arguments, location),
4858
"module_name" => module_name(interner, arguments, location),
@@ -737,6 +747,176 @@ fn function_def_return_type(
737747
Ok(Value::Type(func_meta.return_type().follow_bindings()))
738748
}
739749

750+
// fn set_body(self, body: Quoted)
751+
fn function_def_set_body(
752+
interpreter: &mut Interpreter,
753+
arguments: Vec<(Value, Location)>,
754+
location: Location,
755+
) -> IResult<Value> {
756+
let (self_argument, body_argument) = check_two_arguments(arguments, location)?;
757+
let func_id = get_function_def(self_argument, location)?;
758+
759+
let func_meta = interpreter.elaborator.interner.function_meta(&func_id);
760+
match func_meta.function_body {
761+
FunctionBody::Unresolved(_, _, _) => (),
762+
FunctionBody::Resolving | FunctionBody::Resolved => {
763+
return Err(InterpreterError::CannotMutateFunction { location })
764+
}
765+
}
766+
767+
let body_tokens = get_quoted(body_argument, location)?;
768+
let mut body_quoted = add_token_spans(body_tokens.clone(), location.span);
769+
770+
// Surround the body in `{ ... }` so we can parse it as a block
771+
body_quoted.0.insert(0, SpannedToken::new(Token::LeftBrace, location.span));
772+
body_quoted.0.push(SpannedToken::new(Token::RightBrace, location.span));
773+
774+
let body =
775+
parser::block(parser::fresh_statement()).parse(body_quoted).map_err(|mut errors| {
776+
let error = errors.swap_remove(0);
777+
let rule = "a block";
778+
InterpreterError::FailedToParseMacro {
779+
error,
780+
tokens: body_tokens,
781+
rule,
782+
file: location.file,
783+
}
784+
})?;
785+
786+
let func_meta = interpreter.elaborator.interner.function_meta_mut(&func_id);
787+
func_meta.has_body = true;
788+
func_meta.function_body = FunctionBody::Unresolved(FunctionKind::Normal, body, location.span);
789+
790+
Ok(Value::Unit)
791+
}
792+
793+
// fn set_parameters(self, parameters: [(Quoted, Type)])
794+
fn function_def_set_parameters(
795+
interpreter: &mut Interpreter,
796+
arguments: Vec<(Value, Location)>,
797+
location: Location,
798+
) -> IResult<Value> {
799+
let (self_argument, parameters_argument) = check_two_arguments(arguments, location)?;
800+
801+
let func_id = get_function_def(self_argument, location)?;
802+
let func_meta = interpreter.elaborator.interner.function_meta(&func_id);
803+
match func_meta.function_body {
804+
FunctionBody::Unresolved(_, _, _) => (),
805+
FunctionBody::Resolving | FunctionBody::Resolved => {
806+
return Err(InterpreterError::CannotMutateFunction { location })
807+
}
808+
}
809+
810+
let (input_parameters, _type) =
811+
get_slice(interpreter.elaborator.interner, parameters_argument, location)?;
812+
813+
// What follows is very similar to what happens in Elaborator::define_function_meta
814+
let mut parameters = Vec::new();
815+
let mut parameter_types = Vec::new();
816+
let mut parameter_idents = Vec::new();
817+
818+
for input_parameter in input_parameters {
819+
let mut tuple = get_tuple(interpreter.elaborator.interner, input_parameter, location)?;
820+
let parameter_type = get_type(tuple.pop().unwrap(), location)?;
821+
let parameter_name_tokens = get_quoted(tuple.pop().unwrap(), location)?;
822+
let parameter_name_quoted = add_token_spans(parameter_name_tokens.clone(), location.span);
823+
let parameter_pattern =
824+
parser::pattern().parse(parameter_name_quoted).map_err(|mut errors| {
825+
let error = errors.swap_remove(0);
826+
let rule = "a pattern";
827+
InterpreterError::FailedToParseMacro {
828+
error,
829+
tokens: parameter_name_tokens,
830+
rule,
831+
file: location.file,
832+
}
833+
})?;
834+
835+
let hir_pattern = interpreter.elaborate_item(interpreter.current_function, |elaborator| {
836+
elaborator.elaborate_pattern_and_store_ids(
837+
parameter_pattern,
838+
parameter_type.clone(),
839+
DefinitionKind::Local(None),
840+
&mut parameter_idents,
841+
None,
842+
)
843+
});
844+
845+
parameters.push((hir_pattern, parameter_type.clone(), Visibility::Private));
846+
parameter_types.push(parameter_type);
847+
}
848+
849+
let func_meta = interpreter.elaborator.interner.function_meta(&func_id);
850+
let function_type = Type::Function(
851+
parameter_types,
852+
Box::new(func_meta.return_type().clone()),
853+
Box::new(Type::Unit),
854+
);
855+
856+
// TODO: generics. In Elaborator::define_function_meta there's this code:
857+
//
858+
// if !generics.is_empty() {
859+
// typ = Type::Forall(generics, Box::new(typ));
860+
// }
861+
862+
interpreter.elaborator.interner.push_definition_type(func_meta.name.id, function_type.clone());
863+
864+
{
865+
let func_meta = interpreter.elaborator.interner.function_meta_mut(&func_id);
866+
func_meta.parameters = parameters.into();
867+
func_meta.parameter_idents = parameter_idents;
868+
func_meta.typ = function_type;
869+
}
870+
871+
Ok(Value::Unit)
872+
}
873+
874+
// fn set_return_type(self, return_type: Type)
875+
fn function_def_set_return_type(
876+
interpreter: &mut Interpreter,
877+
arguments: Vec<(Value, Location)>,
878+
location: Location,
879+
) -> IResult<Value> {
880+
let (self_argument, return_type_argument) = check_two_arguments(arguments, location)?;
881+
let return_type = get_type(return_type_argument, location)?;
882+
883+
let func_id = get_function_def(self_argument, location)?;
884+
let func_meta = interpreter.elaborator.interner.function_meta(&func_id);
885+
match func_meta.function_body {
886+
FunctionBody::Unresolved(_, _, _) => (),
887+
FunctionBody::Resolving | FunctionBody::Resolved => {
888+
return Err(InterpreterError::CannotMutateFunction { location })
889+
}
890+
}
891+
892+
let parameter_types = func_meta.parameters.iter().map(|(_, typ, _)| typ.clone()).collect();
893+
894+
let func_meta = interpreter.elaborator.interner.function_meta(&func_id);
895+
let function_type =
896+
Type::Function(parameter_types, Box::new(return_type.clone()), Box::new(Type::Unit));
897+
898+
// TODO: generics. In Elaborator::define_function_meta there's this code:
899+
//
900+
// if !generics.is_empty() {
901+
// typ = Type::Forall(generics, Box::new(typ));
902+
// }
903+
904+
interpreter.elaborator.interner.push_definition_type(func_meta.name.id, function_type.clone());
905+
906+
let quoted_type_id = interpreter.elaborator.interner.push_quoted_type(return_type);
907+
908+
{
909+
let func_meta = interpreter.elaborator.interner.function_meta_mut(&func_id);
910+
func_meta.return_type = FunctionReturnType::Ty(UnresolvedType {
911+
typ: UnresolvedTypeData::Resolved(quoted_type_id),
912+
span: Some(location.span),
913+
});
914+
func_meta.typ = function_type;
915+
}
916+
917+
Ok(Value::Unit)
918+
}
919+
740920
// fn functions(self) -> [FunctionDefinition]
741921
fn module_functions(
742922
interpreter: &Interpreter,

‎compiler/noirc_frontend/src/hir/comptime/interpreter/builtin/builtin_helpers.rs‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -95,6 +95,22 @@ pub(crate) fn get_slice(
9595
}
9696
}
9797

98+
pub(crate) fn get_tuple(
99+
interner: &NodeInterner,
100+
value: Value,
101+
location: Location,
102+
) -> IResult<Vec<Value>> {
103+
match value {
104+
Value::Tuple(values) => Ok(values),
105+
value => {
106+
let type_var = interner.next_type_variable();
107+
let expected = Type::Tuple(vec![type_var]);
108+
let actual = value.get_type().into_owned();
109+
Err(InterpreterError::TypeMismatch { expected, actual, location })
110+
}
111+
}
112+
}
113+
98114
pub(crate) fn get_field(value: Value, location: Location) -> IResult<FieldElement> {
99115
match value {
100116
Value::Field(value) => Ok(value),

‎compiler/noirc_frontend/src/node_interner.rs‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -980,6 +980,10 @@ impl NodeInterner {
980980
self.func_meta.get(func_id).expect("ice: all function ids should have metadata")
981981
}
982982

983+
pub fn function_meta_mut(&mut self, func_id: &FuncId) -> &mut FuncMeta {
984+
self.func_meta.get_mut(func_id).expect("ice: all function ids should have metadata")
985+
}
986+
983987
pub fn try_function_meta(&self, func_id: &FuncId) -> Option<&FuncMeta> {
984988
self.func_meta.get(func_id)
985989
}

‎compiler/noirc_frontend/src/parser/mod.rs‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,8 @@ pub use errors::ParserError;
2323
pub use errors::ParserErrorReason;
2424
use noirc_errors::Span;
2525
pub use parser::{
26-
expression, parse_program, parse_type, path_no_turbofish, top_level_items, trait_bound,
26+
block, expression, fresh_statement, parse_program, parse_type, path_no_turbofish, pattern,
27+
top_level_items, trait_bound,
2728
};
2829

2930
#[derive(Debug, Clone)]

‎compiler/noirc_frontend/src/parser/parser.rs‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -380,7 +380,7 @@ fn block_expr<'a>(
380380
block(statement).map(ExpressionKind::Block).map_with_span(Expression::new)
381381
}
382382

383-
fn block<'a>(
383+
pub fn block<'a>(
384384
statement: impl NoirParser<StatementKind> + 'a,
385385
) -> impl NoirParser<BlockExpression> + 'a {
386386
use Token::*;
@@ -482,7 +482,7 @@ where
482482
})
483483
}
484484

485-
fn fresh_statement() -> impl NoirParser<StatementKind> {
485+
pub fn fresh_statement() -> impl NoirParser<StatementKind> {
486486
statement(expression(), expression_no_constructors(expression()))
487487
}
488488

@@ -538,7 +538,7 @@ where
538538
p.map(StatementKind::new_let)
539539
}
540540

541-
fn pattern() -> impl NoirParser<Pattern> {
541+
pub fn pattern() -> impl NoirParser<Pattern> {
542542
recursive(|pattern| {
543543
let ident_pattern = ident().map(Pattern::Identifier).map_err(|mut error| {
544544
if matches!(error.found(), Token::IntType(..)) {

‎noir_stdlib/src/meta/function_def.nr‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,4 +7,13 @@ impl FunctionDefinition {
77

88
#[builtin(function_def_return_type)]
99
fn return_type(self) -> Type {}
10+
11+
#[builtin(function_def_set_body)]
12+
fn set_body(mut self, body: Quoted) {}
13+
14+
#[builtin(function_def_set_parameters)]
15+
fn set_parameters(mut self, parameters: [(Quoted, Type)]) {}
16+
17+
#[builtin(function_def_set_return_type)]
18+
fn set_return_type(mut self, return_type: Type) {}
1019
}

‎test_programs/compile_success_empty/comptime_function_definition/src/main.nr‎

Lines changed: 22 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,9 +4,13 @@ struct Foo { x: Field, field: Field }
44

55
#[function_attr]
66
fn foo(w: i32, y: Field, Foo { x, field: some_field }: Foo, mut a: bool, (b, c): (i32, i32)) -> i32 {
7+
let _ = (w, y, x, some_field, a, b, c);
78
1
89
}
910

11+
#[mutate_add_one]
12+
fn add_one() {}
13+
1014
comptime fn function_attr(f: FunctionDefinition) {
1115
// Check FunctionDefinition::parameters
1216
let parameters = f.parameters();
@@ -33,4 +37,21 @@ comptime fn function_attr(f: FunctionDefinition) {
3337
assert_eq(f.name(), quote { foo });
3438
}
3539

36-
fn main() {}
40+
comptime fn mutate_add_one(f: FunctionDefinition) {
41+
// fn add_one(x: Field)
42+
assert_eq(f.parameters().len(), 0);
43+
f.set_parameters(&[(quote { x }, type_of(0))]);
44+
assert_eq(f.parameters().len(), 1);
45+
46+
// fn add_one(x: Field) -> Field
47+
assert_eq(f.return_type(), type_of(()));
48+
f.set_return_type(type_of(0));
49+
assert_eq(f.return_type(), type_of(0));
50+
51+
// fn add_one(x: Field) -> Field { x + 1 }
52+
f.set_body(quote { x + 1 });
53+
}
54+
55+
fn main() {
56+
assert_eq(add_one(41), 42);
57+
}

0 commit comments

Comments
 (0)