Skip to content

Commit

Permalink
fix(frontend): Correctly monomorphize turbofish functions (#5049)
Browse files Browse the repository at this point in the history
# Description

## Problem\*

Resolves <!-- Link to GitHub Issue -->

Working towards #4514 and fully
resolving #4710

## Summary\*

We need to accurately monomorphize functions with turbofish operators
where no function parameters or the return type using the specified
generic.

Without this change the following would pass:
```
fn main(x: Field, y: pub Field) {
    let mut hasher = PoseidonHasher::default();
    hasher.write(x);
    hasher.write(y);
    let poseidon_expected_hash = hasher.finish();
    assert(hash_simple_array::<PoseidonHasher>([x, y]) == poseidon_expected_hash);

    let mut hasher = Poseidon2Hasher::default();
    hasher.write(x);
    hasher.write(y);
    let poseidon2_expected_hash = hasher.finish();
    assert(hash_simple_array::<Poseidon2Hasher>([x, y]) == poseidon_expected_hash);
    assert(hash_simple_array::<Poseidon2Hasher>([x, y]) != poseidon2_expected_hash);
}

fn hash_simple_array<H>(input: [Field; 2]) -> Field where H: Hasher + Default {
    let mut hasher: H = H::default();
    hasher.write(input[0]);
    hasher.write(input[1]);
    hasher.finish()
}
```

Essentially any future invocations of `hash_simple_array` would use the
`PoseidonHasher`. We now correctly monomorphize the functions to use the
correct type.

## Additional Context



## Documentation\*

Check one:
- [ ] No documentation needed.
- [ ] Documentation included in this PR.
- [ ] **[For Experimental Features]** Documentation to be submitted in a
separate PR.

# PR Checklist\*

- [ ] I have tested the changes locally.
- [ ] I have formatted the changes with [Prettier](https://prettier.io/)
and/or `cargo fmt` on default settings.

---------

Co-authored-by: Jake Fecher <jake@aztecprotocol.com>
Co-authored-by: Jake Fecher <jfecher11@gmail.com>
Co-authored-by: Tom French <tom@tomfren.ch>
  • Loading branch information
4 people committed May 21, 2024
1 parent 89846cf commit fd772e7
Show file tree
Hide file tree
Showing 5 changed files with 81 additions and 34 deletions.
48 changes: 36 additions & 12 deletions compiler/noirc_frontend/src/monomorphization/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -55,9 +55,12 @@ struct LambdaContext {
struct Monomorphizer<'interner> {
/// Functions are keyed by their unique ID and expected type so that we can monomorphize
/// a new version of the function for each type.
/// We also key by any turbofish generics that are specified.
/// This is necessary for a case where we may have a trait generic that can be instantiated
/// outside of a function parameter or return value.
///
/// Using nested HashMaps here lets us avoid cloning HirTypes when calling .get()
functions: HashMap<node_interner::FuncId, HashMap<HirType, FuncId>>,
functions: HashMap<node_interner::FuncId, HashMap<(HirType, Vec<HirType>), FuncId>>,

/// Unlike functions, locals are only keyed by their unique ID because they are never
/// duplicated during monomorphization. Doing so would allow them to be used polymorphically
Expand Down Expand Up @@ -198,10 +201,15 @@ impl<'interner> Monomorphizer<'interner> {
id: node_interner::FuncId,
expr_id: node_interner::ExprId,
typ: &HirType,
turbofish_generics: Vec<HirType>,
trait_method: Option<TraitMethodId>,
) -> Definition {
let typ = typ.follow_bindings();
match self.functions.get(&id).and_then(|inner_map| inner_map.get(&typ)) {
match self
.functions
.get(&id)
.and_then(|inner_map| inner_map.get(&(typ.clone(), turbofish_generics.clone())))
{
Some(id) => Definition::Function(*id),
None => {
// Function has not been monomorphized yet
Expand All @@ -222,7 +230,8 @@ impl<'interner> Monomorphizer<'interner> {
Definition::Builtin(opcode)
}
FunctionKind::Normal => {
let id = self.queue_function(id, expr_id, typ, trait_method);
let id =
self.queue_function(id, expr_id, typ, turbofish_generics, trait_method);
Definition::Function(id)
}
FunctionKind::Oracle => {
Expand All @@ -249,8 +258,14 @@ impl<'interner> Monomorphizer<'interner> {
}

/// Prerequisite: typ = typ.follow_bindings()
fn define_function(&mut self, id: node_interner::FuncId, typ: HirType, new_id: FuncId) {
self.functions.entry(id).or_default().insert(typ, new_id);
fn define_function(
&mut self,
id: node_interner::FuncId,
typ: HirType,
turbofish_generics: Vec<HirType>,
new_id: FuncId,
) {
self.functions.entry(id).or_default().insert((typ, turbofish_generics), new_id);
}

fn compile_main(
Expand Down Expand Up @@ -393,7 +408,7 @@ impl<'interner> Monomorphizer<'interner> {
use ast::Literal::*;

let expr = match self.interner.expression(&expr) {
HirExpression::Ident(ident, _) => self.ident(ident, expr)?,
HirExpression::Ident(ident, generics) => self.ident(ident, expr, generics)?,
HirExpression::Literal(HirLiteral::Str(contents)) => Literal(Str(contents)),
HirExpression::Literal(HirLiteral::FmtStr(contents, idents)) => {
let fields = try_vecmap(idents, |ident| self.expr(ident))?;
Expand Down Expand Up @@ -825,6 +840,7 @@ impl<'interner> Monomorphizer<'interner> {
&mut self,
ident: HirIdent,
expr_id: node_interner::ExprId,
generics: Option<Vec<HirType>>,
) -> Result<ast::Expression, MonomorphizationError> {
let typ = self.interner.id_type(expr_id);

Expand All @@ -838,7 +854,13 @@ impl<'interner> Monomorphizer<'interner> {
let mutable = definition.mutable;
let location = Some(ident.location);
let name = definition.name.clone();
let definition = self.lookup_function(*func_id, expr_id, &typ, None);
let definition = self.lookup_function(
*func_id,
expr_id,
&typ,
generics.unwrap_or_default(),
None,
);
let typ = Self::convert_type(&typ, ident.location)?;
let ident = ast::Ident { location, mutable, definition, name, typ: typ.clone() };
let ident_expression = ast::Expression::Ident(ident);
Expand Down Expand Up @@ -1063,10 +1085,11 @@ impl<'interner> Monomorphizer<'interner> {
}
};

let func_id = match self.lookup_function(func_id, expr_id, &function_type, Some(method)) {
Definition::Function(func_id) => func_id,
_ => unreachable!(),
};
let func_id =
match self.lookup_function(func_id, expr_id, &function_type, vec![], Some(method)) {
Definition::Function(func_id) => func_id,
_ => unreachable!(),
};

let the_trait = self.interner.get_trait(method.trait_id);
let location = self.interner.expr_location(&expr_id);
Expand Down Expand Up @@ -1292,10 +1315,11 @@ impl<'interner> Monomorphizer<'interner> {
id: node_interner::FuncId,
expr_id: node_interner::ExprId,
function_type: HirType,
turbofish_generics: Vec<HirType>,
trait_method: Option<TraitMethodId>,
) -> FuncId {
let new_id = self.next_function_id();
self.define_function(id, function_type.clone(), new_id);
self.define_function(id, function_type.clone(), turbofish_generics, new_id);

let bindings = self.interner.get_instantiation_bindings(expr_id);
let bindings = self.follow_bindings(bindings);
Expand Down
22 changes: 0 additions & 22 deletions test_programs/execution_success/trait_method_mut_self/src/main.nr
Original file line number Diff line number Diff line change
Expand Up @@ -12,14 +12,6 @@ fn main(x: Field, y: pub Field) {

pass_trait_by_mut_ref(&mut a_mut_ref, y);
assert(a_mut_ref.x == y);

let mut hasher = Poseidon2Hasher::default();
hasher.write(x);
hasher.write(y);
let expected_hash = hasher.finish();
// Check that we get the same result when using the hasher in a
// method that purely uses trait methods without a supplied implementation.
assert(hash_simple_array::<Poseidon2Hasher>([x, y]) == expected_hash);
}

trait SomeTrait {
Expand Down Expand Up @@ -58,17 +50,3 @@ fn pass_trait_by_mut_ref<T>(a_mut_ref: &mut T, value: Field) where T: SomeTrait
// We auto add a mutable reference to the object type if the method call expects a mutable self
a_mut_ref.set_value(value);
}

fn hash_simple_array<H>(input: [Field; 2]) -> Field where H: Hasher + Default {
// Check that we can call a trait method instead of a trait implementation
// TODO: Need to remove the need for this type annotation
// TODO: Curently, without the annotation we will get `Expression type is ambiguous` when trying to use the `hasher`
let mut hasher: H = H::default();
// Regression that the object is converted to a mutable reference type `&mut _`.
// Otherwise will see `Expected type &mut _, found type H`.
// Then we need to make sure to also auto dereference later in the type checking process
// when searching for a matching impl or else we will get `No matching impl found for `&mut H: Hasher`
hasher.write(input[0]);
hasher.write(input[1]);
hasher.finish()
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,7 @@
[package]
name = "turbofish_call_func_diff_types"
type = "bin"
authors = [""]
compiler_version = ">=0.29.0"

[dependencies]
Original file line number Diff line number Diff line change
@@ -0,0 +1,2 @@
x = "5"
y = "10"
Original file line number Diff line number Diff line change
@@ -0,0 +1,36 @@
use dep::std::hash::Hasher;
use dep::std::hash::poseidon2::Poseidon2Hasher;
use dep::std::hash::poseidon::PoseidonHasher;

fn main(x: Field, y: pub Field) {
let mut hasher = PoseidonHasher::default();
hasher.write(x);
hasher.write(y);
let poseidon_expected_hash = hasher.finish();
// Check that we get the same result when using the hasher in a
// method that purely uses trait methods without a supplied implementation.
assert(hash_simple_array::<PoseidonHasher>([x, y]) == poseidon_expected_hash);

// Now let's do the same logic but with a different `Hasher` supplied to the turbofish operator
// We want to make sure that we have correctly monomorphized a function with a trait generic
// where the generic is not used on any function parameters or the return value.
let mut hasher = Poseidon2Hasher::default();
hasher.write(x);
hasher.write(y);
let poseidon2_expected_hash = hasher.finish();
assert(hash_simple_array::<Poseidon2Hasher>([x, y]) == poseidon2_expected_hash);
}

fn hash_simple_array<H>(input: [Field; 2]) -> Field where H: Hasher + Default {
// Check that we can call a trait method instead of a trait implementation
// TODO(https://github.com/noir-lang/noir/issues/5063): Need to remove the need for this type annotation
// Curently, without the annotation we will get `Expression type is ambiguous` when trying to use the `hasher`
let mut hasher: H = H::default();
// Regression that the object is converted to a mutable reference type `&mut _`.
// Otherwise will see `Expected type &mut _, found type H`.
// Then we need to make sure to also auto dereference later in the type checking process
// when searching for a matching impl or else we will get `No matching impl found for `&mut H: Hasher`
hasher.write(input[0]);
hasher.write(input[1]);
hasher.finish()
}

0 comments on commit fd772e7

Please sign in to comment.