Skip to content

Commit

Permalink
fix: binary compute_signature returning a PolyFuncType with binders (#…
Browse files Browse the repository at this point in the history
…710)

Add test that we can pass typevars to TypeScheme part but not binary
part. Fix inconsistency between validate_args and compute_signature.

Also add `impl From<Type> for TypeArg`, this is so much easier than the
awkward brace syntax `TypeArg::Type {ty:...}`, tho sometimes that's good
for explicitness :)
  • Loading branch information
acl-cqc authored Nov 22, 2023
1 parent 2f952ab commit b7a5176
Show file tree
Hide file tree
Showing 2 changed files with 127 additions and 4 deletions.
125 changes: 121 additions & 4 deletions src/extension/op_def.rs
Original file line number Diff line number Diff line change
Expand Up @@ -180,9 +180,6 @@ impl OpDef {
args: &[TypeArg],
exts: &ExtensionRegistry,
) -> Result<FunctionType, SignatureError> {
// Hugr's are monomorphic, so check the args have no free variables
args.iter().try_for_each(|ta| ta.validate(exts, &[]))?;

let temp: PolyFuncType; // to keep alive
let (pf, args) = match &self.signature_func {
SignatureFunc::TypeScheme(ts) => (ts, args),
Expand Down Expand Up @@ -369,11 +366,16 @@ impl Extension {

#[cfg(test)]
mod test {
use std::num::NonZeroU64;

use smol_str::SmolStr;

use crate::builder::{DFGBuilder, Dataflow, DataflowHugr};
use crate::extension::prelude::USIZE_T;
use crate::extension::{ExtensionRegistry, PRELUDE};
use crate::extension::{
CustomSignatureFunc, ExtensionRegistry, SignatureError, EMPTY_REG, PRELUDE,
PRELUDE_REGISTRY,
};
use crate::ops::custom::ExternalOp;
use crate::ops::LeafOp;
use crate::std_extensions::collections::{EXTENSION, LIST_TYPENAME};
Expand Down Expand Up @@ -413,4 +415,119 @@ mod test {

Ok(())
}

#[test]
fn binary_polyfunc() -> Result<(), Box<dyn std::error::Error>> {
// Test a custom binary `compute_signature` that returns a PolyFuncType
// where the latter declares more type params itself. In particular,
// we should be able to substitute (external) type variables into the latter,
// but not pass them into the former (custom binary function).
struct SigFun();
impl CustomSignatureFunc for SigFun {
fn compute_signature(
&self,
_name: &SmolStr,
arg_values: &[TypeArg],
_misc: &std::collections::HashMap<String, serde_yaml::Value>,
_exts: &ExtensionRegistry,
) -> Result<PolyFuncType, SignatureError> {
const TP: TypeParam = TypeParam::Type(TypeBound::Any);
let [TypeArg::BoundedNat { n }] = arg_values else {
return Err(SignatureError::InvalidTypeArgs);
};
let n = *n as usize;
let tvs: Vec<Type> = (0..n)
.map(|_| Type::new_var_use(0, TypeBound::Any))
.collect();
Ok(PolyFuncType::new(
vec![TP],
FunctionType::new(tvs.clone(), vec![Type::new_tuple(tvs)]),
))
}
}
let mut e = Extension::new(EXT_ID);
let def = e.add_op_custom_sig_simple(
"MyOp".into(),
"".to_string(),
vec![TypeParam::max_nat()],
SigFun(),
)?;

// Base case, no type variables:
let args = [TypeArg::BoundedNat { n: 3 }, USIZE_T.into()];
assert_eq!(
def.compute_signature(&args, &PRELUDE_REGISTRY),
Ok(FunctionType::new(
vec![USIZE_T; 3],
vec![Type::new_tuple(vec![USIZE_T; 3])]
))
);
assert_eq!(def.validate_args(&args, &PRELUDE_REGISTRY, &[]), Ok(()));

// Second arg may be a variable (substitutable)
let tyvar = Type::new_var_use(0, TypeBound::Eq);
let tyvars: Vec<Type> = vec![tyvar.clone(); 3];
let args = [TypeArg::BoundedNat { n: 3 }, tyvar.clone().into()];
assert_eq!(
def.compute_signature(&args, &PRELUDE_REGISTRY),
Ok(FunctionType::new(
tyvars.clone(),
vec![Type::new_tuple(tyvars)]
))
);
def.validate_args(&args, &PRELUDE_REGISTRY, &[TypeParam::Type(TypeBound::Eq)])
.unwrap();

// quick sanity check that we are validating the args - note changed bound:
assert_eq!(
def.validate_args(&args, &PRELUDE_REGISTRY, &[TypeParam::Type(TypeBound::Any)]),
Err(SignatureError::TypeVarDoesNotMatchDeclaration {
actual: TypeBound::Any.into(),
cached: TypeBound::Eq.into()
})
);

// First arg must be concrete, not a variable
let kind = TypeParam::bounded_nat(NonZeroU64::new(5).unwrap());
let args = [TypeArg::new_var_use(0, kind.clone()), USIZE_T.into()];
// We can't prevent this from getting into our compute_signature implementation:
assert_eq!(
def.compute_signature(&args, &PRELUDE_REGISTRY),
Err(SignatureError::InvalidTypeArgs)
);
// But validation rules it out, even when the variable is declared:
assert_eq!(
def.validate_args(&args, &PRELUDE_REGISTRY, &[kind]),
Err(SignatureError::FreeTypeVar {
idx: 0,
num_decls: 0
})
);

Ok(())
}

#[test]
fn type_scheme_instantiate_var() -> Result<(), Box<dyn std::error::Error>> {
// Check that we can instantiate a PolyFuncType-scheme with an (external)
// type variable
let mut e = Extension::new(EXT_ID);
let def = e.add_op_type_scheme_simple(
"SimpleOp".into(),
"".into(),
PolyFuncType::new(
vec![TypeParam::Type(TypeBound::Any)],
FunctionType::new_endo(vec![Type::new_var_use(0, TypeBound::Any)]),
),
)?;
let tv = Type::new_var_use(1, TypeBound::Eq);
let args = [TypeArg::Type { ty: tv.clone() }];
let decls = [TypeParam::Extensions, TypeBound::Eq.into()];
def.validate_args(&args, &EMPTY_REG, &decls).unwrap();
assert_eq!(
def.compute_signature(&args, &EMPTY_REG),
Ok(FunctionType::new_endo(vec![tv]))
);
Ok(())
}
}
6 changes: 6 additions & 0 deletions src/types/type_param.rs
Original file line number Diff line number Diff line change
Expand Up @@ -97,6 +97,12 @@ impl From<TypeBound> for TypeParam {
}
}

impl From<Type> for TypeArg {
fn from(ty: Type) -> Self {
Self::Type { ty }
}
}

/// A statically-known argument value to an operation.
#[derive(Clone, Debug, PartialEq, Eq, serde::Deserialize, serde::Serialize)]
#[non_exhaustive]
Expand Down

0 comments on commit b7a5176

Please sign in to comment.