Repository: qir-alliance/pyqir (language: python)

Issue:
Remove TypeFactory
With the type constructors added in #180, there are now two ways to make types. We should remove TypeFactory and make the normal constructors the only way. A context property needs to be added to SimpleModule so it can be passed to the type constructors when adding external functions to SimpleModules.


Discussion / hints:

Fix patch:
diff --git a/examples/arithmetic.py b/examples/arithmetic.py
index 3f5a0986..a99fd2fb 100644
--- a/examples/arithmetic.py
+++ b/examples/arithmetic.py
@@ -1,31 +1,32 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
 
-from pyqir import BasicQisBuilder, SimpleModule, const
+import pyqir
 
-mod = SimpleModule("arithmetic", num_qubits=0, num_results=0)
-qis = BasicQisBuilder(mod.builder)
-types = mod.types
+mod = pyqir.SimpleModule("arithmetic", num_qubits=0, num_results=0)
+qis = pyqir.BasicQisBuilder(mod.builder)
 
 # Declare functions that can produce and consume integers at runtime. See
 # external_functions.py.
-i32 = types.int(32)
-get_int = mod.add_external_function("get_int", types.function(i32, []))
-take_int = mod.add_external_function("take_int", types.function(types.void, [i32]))
+i32 = pyqir.IntType(mod.context, 32)
+get_int = mod.add_external_function("get_int", pyqir.FunctionType(i32, []))
+take_int = mod.add_external_function(
+    "take_int", pyqir.FunctionType(pyqir.Type.void(mod.context), [i32])
+)
 
 # Add 3 to a number and multiply the result by 2.
 a = mod.builder.call(get_int, [])
 assert a is not None
 # Python numbers need to be converted into QIR constant values. Since it's being
 # added to a 32-bit integer returned by get_int, its type needs to be the same.
-three = const(i32, 3)
+three = pyqir.const(i32, 3)
 b = mod.builder.add(three, a)
-c = mod.builder.mul(const(i32, 2), b)
+c = mod.builder.mul(pyqir.const(i32, 2), b)
 
 # Negation can be done by subtracting an integer from zero.
 x = mod.builder.call(get_int, [])
 assert x is not None
-negative_x = mod.builder.sub(const(i32, 0), x)
+negative_x = mod.builder.sub(pyqir.const(i32, 0), x)
 
 # Consume the results.
 mod.builder.call(take_int, [c])
diff --git a/examples/dynamic_allocation.py b/examples/dynamic_allocation.py
index 1d4e738b..b96aff33 100644
--- a/examples/dynamic_allocation.py
+++ b/examples/dynamic_allocation.py
@@ -1,28 +1,31 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
 
-from pyqir import BasicQisBuilder, SimpleModule
+import pyqir
 
-# PyQIR Generator assumes you want to use static allocation for qubits and
-# results, but you can still use dynamic allocation by manually calling the
-# appropriate runtime functions.
-mod = SimpleModule("dynamic_allocation", num_qubits=0, num_results=0)
-types = mod.types
+mod = pyqir.SimpleModule("dynamic_allocation", num_qubits=0, num_results=0)
+qubit_type = pyqir.qubit_type(mod.context)
+result_type = pyqir.result_type(mod.context)
+
+# PyQIR assumes you want to use static allocation for qubits and results, but
+# you can still use dynamic allocation by manually calling the appropriate
+# runtime functions.
 qubit_allocate = mod.add_external_function(
-    "__quantum__rt__qubit_allocate", types.function(types.qubit, [])
+    "__quantum__rt__qubit_allocate", pyqir.FunctionType(qubit_type, [])
 )
 qubit_release = mod.add_external_function(
-    "__quantum__rt__qubit_release", types.function(types.void, [types.qubit])
+    "__quantum__rt__qubit_release",
+    pyqir.FunctionType(pyqir.Type.void(mod.context), [qubit_type]),
 )
 result_get_one = mod.add_external_function(
-    "__quantum__rt__result_get_one", types.function(types.result, [])
+    "__quantum__rt__result_get_one", pyqir.FunctionType(result_type, [])
 )
 result_equal = mod.add_external_function(
     "__quantum__rt__result_equal",
-    types.function(types.bool, [types.result, types.result]),
+    pyqir.FunctionType(pyqir.IntType(mod.context, 1), [result_type, result_type]),
 )
 m = mod.add_external_function(
-    "__quantum__qis__m__body", types.function(types.result, [types.qubit])
+    "__quantum__qis__m__body", pyqir.FunctionType(result_type, [qubit_type])
 )
 
 # Instead of mod.qubits[i], use __quantum__rt__qubit_allocate.
@@ -30,7 +33,7 @@
 assert qubit_return is not None
 qubit = qubit_return
 
-qis = BasicQisBuilder(mod.builder)
+qis = pyqir.BasicQisBuilder(mod.builder)
 qis.h(qubit)
 
 # Instead of qis.mz, use __quantum__qis__m__body.
diff --git a/examples/external_functions.py b/examples/external_functions.py
index 67242042..6b65f90b 100644
--- a/examples/external_functions.py
+++ b/examples/external_functions.py
@@ -1,29 +1,33 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
 
-from pyqir import BasicQisBuilder, SimpleModule, const
+import pyqir
 
-mod = SimpleModule("external_functions", num_qubits=1, num_results=0)
-qis = BasicQisBuilder(mod.builder)
-types = mod.types
+mod = pyqir.SimpleModule("external_functions", num_qubits=1, num_results=0)
+qis = pyqir.BasicQisBuilder(mod.builder)
+
+void = pyqir.Type.void(mod.context)
+i64 = pyqir.IntType(mod.context, 64)
 
 # Declare an externally linked function named that has no parameters and returns
 # void.
-my_function = mod.add_external_function("my_function", types.function(types.void, []))
+my_function = mod.add_external_function("my_function", pyqir.FunctionType(void, []))
 
 # Call the function with no arguments (the empty list).
 mod.builder.call(my_function, [])
 
 # Declare a function that takes an integer and a qubit and returns void.
 my_gate = mod.add_external_function(
-    "my_gate", types.function(types.void, [types.int(64), types.qubit])
+    "my_gate", pyqir.FunctionType(void, [i64, pyqir.qubit_type(mod.context)])
 )
 
 # Call the functions with a list of arguments.
-mod.builder.call(my_gate, [const(types.int(64), 123), mod.qubits[0]])
+mod.builder.call(my_gate, [pyqir.const(i64, 123), mod.qubits[0]])
 
 # Declare a function that returns a double.
-get_angle = mod.add_external_function("get_angle", types.function(types.double, []))
+get_angle = mod.add_external_function(
+    "get_angle", pyqir.FunctionType(pyqir.Type.double(mod.context), [])
+)
 
 # Use the return value of the function as the input to a rotation gate.
 angle = mod.builder.call(get_angle, [])
diff --git a/examples/if_bool.py b/examples/if_bool.py
index 4e15b908..ed56ba8e 100644
--- a/examples/if_bool.py
+++ b/examples/if_bool.py
@@ -1,29 +1,29 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
 
-from pyqir import BasicQisBuilder, IntPredicate, SimpleModule, const
+import pyqir
 
-mod = SimpleModule("if_bool", num_qubits=2, num_results=2)
-qis = BasicQisBuilder(mod.builder)
+mod = pyqir.SimpleModule("if_bool", num_qubits=2, num_results=2)
+qis = pyqir.BasicQisBuilder(mod.builder)
 
 # Use an external function to generate integers that we can compare with icmp.
-i32 = mod.types.int(32)
-get_int = mod.add_external_function("get_int", mod.types.function(i32, []))
+i32 = pyqir.IntType(mod.context, 32)
+get_int = mod.add_external_function("get_int", pyqir.FunctionType(i32, []))
 
 # Apply X to the qubit if 'a' is 7.
 a = mod.builder.call(get_int, [])
 assert a is not None
-a_eq_7 = mod.builder.icmp(IntPredicate.EQ, a, const(i32, 7))
+a_eq_7 = mod.builder.icmp(pyqir.IntPredicate.EQ, a, pyqir.const(i32, 7))
 mod.builder.if_(a_eq_7, lambda: qis.x(mod.qubits[0]))
 
 # Multiple conditions can be combined with 'and' and 'or'.
 b = mod.builder.call(get_int, [])
 assert b is not None
-b_sgt_a = mod.builder.icmp(IntPredicate.SGT, b, a)
+b_sgt_a = mod.builder.icmp(pyqir.IntPredicate.SGT, b, a)
 or_cond = mod.builder.or_(a_eq_7, b_sgt_a)
 
 # Both the true and false branches can be specified.
-b_ne_2 = mod.builder.icmp(IntPredicate.NE, b, const(i32, 2))
+b_ne_2 = mod.builder.icmp(pyqir.IntPredicate.NE, b, pyqir.const(i32, 2))
 and_cond = mod.builder.and_(or_cond, b_ne_2)
 mod.builder.if_(
     and_cond,
diff --git a/pyqir/pyqir/__init__.py b/pyqir/pyqir/__init__.py
index f74d9ad8..c647c774 100644
--- a/pyqir/pyqir/__init__.py
+++ b/pyqir/pyqir/__init__.py
@@ -29,7 +29,6 @@
     StructType,
     Switch,
     Type,
-    TypeFactory,
     Value,
     const_getelementptr,
     const,
@@ -78,7 +77,6 @@
     "StructType",
     "Switch",
     "Type",
-    "TypeFactory",
     "Value",
     "const_getelementptr",
     "const",
diff --git a/pyqir/pyqir/_native.pyi b/pyqir/pyqir/_native.pyi
index e68de9a0..53594e5f 100644
--- a/pyqir/pyqir/_native.pyi
+++ b/pyqir/pyqir/_native.pyi
@@ -8,23 +8,6 @@ from typing import Callable, List, Optional, Sequence, Tuple, Union
 class Context:
     def __init__(self) -> None: ...
 
-class TypeFactory:
-    @property
-    def void(self) -> Type: ...
-    @property
-    def bool(self) -> Type: ...
-    def int(self, width: int) -> Type: ...
-    @property
-    def double(self) -> Type: ...
-    @property
-    def qubit(self) -> Type: ...
-    @property
-    def result(self) -> Type: ...
-    @staticmethod
-    def pointer_to(ty: Type) -> Type: ...
-    @staticmethod
-    def function(ret: Type, params: Sequence[Type]) -> Type: ...
-
 class Builder:
     def __init__(self, context: Context) -> None: ...
     def insert_at_end(self, block: BasicBlock) -> None: ...
@@ -75,9 +58,15 @@ class BasicQisBuilder:
     ) -> None: ...
 
 class SimpleModule:
-    def __init__(self, name: str, num_qubits: int, num_results: int) -> None: ...
+    def __init__(
+        self,
+        name: str,
+        num_qubits: int,
+        num_results: int,
+        context: Optional[Context] = None,
+    ) -> None: ...
     @property
-    def types(self) -> TypeFactory: ...
+    def context(self) -> Context: ...
     @property
     def qubits(self) -> List[Value]: ...
     @property
@@ -86,7 +75,7 @@ class SimpleModule:
     def builder(self) -> Builder: ...
     def ir(self) -> str: ...
     def bitcode(self) -> bytes: ...
-    def add_external_function(self, name: str, ty: Type) -> Function: ...
+    def add_external_function(self, name: str, ty: FunctionType) -> Function: ...
     def add_global_string(self, value: bytes) -> Value: ...
 
 class Type:
@@ -139,7 +128,7 @@ class Value:
     @property
     def type(self) -> Type: ...
     @property
-    def name(self) -> Optional[str]: ...
+    def name(self) -> str: ...
 
 class BasicBlock(Value):
     def __init__(
diff --git a/pyqir/src/python.rs b/pyqir/src/python.rs
index 8aa287af..f5b109f2 100644
--- a/pyqir/src/python.rs
+++ b/pyqir/src/python.rs
@@ -10,7 +10,7 @@ use crate::{
     },
     module::{verify_module, Attribute, Module},
     qis::BasicQisBuilder,
-    simple::{SimpleModule, TypeFactory},
+    simple::SimpleModule,
     types::{
         is_qubit_type, is_result_type, qubit_type, result_type, ArrayType, FunctionType, IntType,
         PointerType, StructType, Type,
@@ -53,7 +53,6 @@ fn _native(_py: Python, m: &PyModule) -> PyResult<()> {
     m.add_class::<StructType>()?;
     m.add_class::<Switch>()?;
     m.add_class::<Type>()?;
-    m.add_class::<TypeFactory>()?;
     m.add_class::<Value>()?;
     m.add_function(wrap_pyfunction!(const_getelementptr, m)?)?;
     m.add_function(wrap_pyfunction!(constant_bytes, m)?)?;
diff --git a/pyqir/src/simple.rs b/pyqir/src/simple.rs
index 6e222d6f..7aa0bb36 100644
--- a/pyqir/src/simple.rs
+++ b/pyqir/src/simple.rs
@@ -5,22 +5,16 @@ use crate::{
     builder::Builder,
     context::{self, Context},
     module::Module,
-    types::Type,
+    types::FunctionType,
     values::Value,
 };
-use inkwell::{
-    context::ContextRef,
-    module::Linkage,
-    types::{AnyTypeEnum, BasicType, BasicTypeEnum},
-    AddressSpace,
-};
+use inkwell::module::Linkage;
 use pyo3::{
     exceptions::{PyOSError, PyUnicodeDecodeError, PyValueError},
     prelude::*,
     types::PyBytes,
 };
-use qirlib::{passes, types, values};
-use std::{convert::Into, mem::transmute};
+use qirlib::{passes, values};
 
 /// A simple module represents an executable program with these restrictions:
 ///
@@ -31,12 +25,12 @@ use std::{convert::Into, mem::transmute};
 /// :param str name: The name of the module.
 /// :param int num_qubits: The number of statically allocated qubits.
 /// :param int num_results: The number of statically allocated results.
+/// :param Optional[Context] context: The global context.
 #[pyclass(unsendable)]
-#[pyo3(text_signature = "(name, num_qubits, num_results)")]
+#[pyo3(text_signature = "(name, num_qubits, num_results, context=None)")]
 pub(crate) struct SimpleModule {
     module: Py<Module>,
     builder: Py<Builder>,
-    types: Py<TypeFactory>,
     num_qubits: u64,
     num_results: u64,
 }
@@ -44,8 +38,14 @@ pub(crate) struct SimpleModule {
 #[pymethods]
 impl SimpleModule {
     #[new]
-    fn new(py: Python, name: &str, num_qubits: u64, num_results: u64) -> PyResult<SimpleModule> {
-        let context = Py::new(py, Context::new())?;
+    fn new(
+        py: Python,
+        name: &str,
+        num_qubits: u64,
+        num_results: u64,
+        context: Option<Py<Context>>,
+    ) -> PyResult<SimpleModule> {
+        let context = context.map_or_else(|| Py::new(py, Context::new()), Ok)?;
         let module = Py::new(py, Module::new(py, context.clone(), name))?;
         let builder = Py::new(py, Builder::new(py, context.clone()))?;
 
@@ -62,15 +62,14 @@ impl SimpleModule {
         Ok(SimpleModule {
             module,
             builder,
-            types: Py::new(py, TypeFactory { context })?,
             num_qubits,
             num_results,
         })
     }
 
     #[getter]
-    fn types(&self) -> Py<TypeFactory> {
-        self.types.clone()
+    fn context(&self, py: Python) -> Py<Context> {
+        self.module.borrow(py).context().clone()
     }
 
     /// The global qubit register.
@@ -134,15 +133,18 @@ impl SimpleModule {
     /// :return: The function value.
     /// :rtype: Function
     #[pyo3(text_signature = "(self, name, ty)")]
-    fn add_external_function(&mut self, py: Python, name: &str, ty: &Type) -> PyResult<PyObject> {
+    fn add_external_function(
+        &mut self,
+        py: Python,
+        name: &str,
+        ty: PyRef<FunctionType>,
+    ) -> PyResult<PyObject> {
         let module = self.module.borrow(py);
+        let function_ty = unsafe { ty.get() };
+        let ty = ty.into_super();
         context::require_same(py, [module.context(), ty.context()])?;
-
-        let context = ty.context().clone();
-        let ty = unsafe { transmute::<AnyTypeEnum<'_>, AnyTypeEnum<'static>>(ty.get()) }
-            .into_function_type();
-        let function = unsafe { module.get() }.add_function(name, ty, None);
-        unsafe { Value::from_any(py, context, function) }
+        let function = unsafe { module.get() }.add_function(name, function_ty, None);
+        unsafe { Value::from_any(py, ty.context().clone(), function) }
     }
 
     /// Adds a global null-terminated string constant to the module.
@@ -185,121 +187,6 @@ impl SimpleModule {
     }
 }
 
-/// Provides access to all supported types.
-#[pyclass]
-pub(crate) struct TypeFactory {
-    context: Py<Context>,
-}
-
-#[pymethods]
-impl TypeFactory {
-    /// The void type.
-    ///
-    /// :type: Type
-    #[getter]
-    fn void(&self, py: Python) -> PyResult<PyObject> {
-        self.new_type(py, |context| context.void_type().into())
-    }
-
-    /// The boolean type.
-    ///
-    /// :type: Type
-    #[getter]
-    fn bool(&self, py: Python) -> PyResult<PyObject> {
-        self.new_type(py, |context| context.bool_type().into())
-    }
-
-    /// An integer type.
-    ///
-    /// :param int width: The number of bits in the integers.
-    /// :returns: The integer type.
-    /// :rtype: Type
-    #[pyo3(text_signature = "(width)")]
-    fn int(&self, py: Python, width: u32) -> PyResult<PyObject> {
-        self.new_type(py, |context| context.custom_width_int_type(width).into())
-    }
-
-    /// The double type.
-    ///
-    /// :type: Type
-    #[getter]
-    fn double(&self, py: Python) -> PyResult<PyObject> {
-        self.new_type(py, |context| context.f64_type().into())
-    }
-
-    /// The qubit type.
-    ///
-    /// :type: Type
-    #[getter]
-    fn qubit(&self, py: Python) -> PyResult<PyObject> {
-        self.new_type(py, |context| types::qubit(context).into())
-    }
-
-    /// The measurement result type.
-    ///
-    /// :type: Type
-    #[getter]
-    fn result(&self, py: Python) -> PyResult<PyObject> {
-        self.new_type(py, |context| types::result(context).into())
-    }
-
-    /// A pointer type.
-    ///
-    /// :param ty: The type pointed to.
-    /// :returns: The pointer type.
-    /// :rtype: Type
-    #[staticmethod]
-    #[pyo3(text_signature = "(ty)")]
-    #[allow(clippy::similar_names)]
-    fn pointer_to(py: Python, ty: &Type) -> PyResult<PyObject> {
-        let pointee = BasicTypeEnum::try_from(unsafe { ty.get() })
-            .map_err(|()| PyValueError::new_err("Type can't be pointed to."))?;
-        let pointer = pointee.ptr_type(AddressSpace::Generic);
-        unsafe { Type::from_any(py, ty.context().clone(), pointer.into()) }
-    }
-
-    /// A function type.
-    ///
-    /// :param Type ret: The return type.
-    /// :param List[Type] params: The parameter types.
-    /// :returns: The function type.
-    /// :rtype: Type
-    #[staticmethod]
-    #[pyo3(text_signature = "(ret, params)")]
-    #[allow(clippy::needless_pass_by_value)]
-    fn function(py: Python, ret: &Type, params: Vec<Py<Type>>) -> PyResult<PyObject> {
-        context::require_same(
-            py,
-            params
-                .iter()
-                .map(|t| t.borrow(py).context().clone())
-                .chain([ret.context().clone()]),
-        )?;
-
-        let ty = crate::types::function(
-            unsafe { &ret.get() },
-            params.iter().map(|t| unsafe {
-                transmute::<AnyTypeEnum<'_>, AnyTypeEnum<'static>>(t.borrow(py).get())
-            }),
-        )
-        .ok_or_else(|| PyValueError::new_err("Invalid return or parameter type."))?;
-
-        unsafe { Type::from_any(py, ret.context().clone(), ty.into()) }
-    }
-}
-
-impl TypeFactory {
-    fn new_type(
-        &self,
-        py: Python,
-        f: impl for<'ctx> Fn(&ContextRef<'ctx>) -> AnyTypeEnum<'ctx>,
-    ) -> PyResult<PyObject> {
-        let context = self.context.borrow(py);
-        let ty = f(&context.void_type().get_context());
-        unsafe { Type::from_any(py, self.context.clone(), ty) }
-    }
-}
-
 fn clone_module<'ctx>(
     module: &inkwell::module::Module,
     context: &'ctx inkwell::context::Context,


Test patch:
diff --git a/pyqir/tests/test_external_functions.py b/pyqir/tests/test_external_functions.py
index 3180564b..2b26a4ce 100644
--- a/pyqir/tests/test_external_functions.py
+++ b/pyqir/tests/test_external_functions.py
@@ -1,33 +1,38 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
 
+import re
+from typing import Any, Callable, List, Union
+
+import pytest
+
+import pyqir
 from pyqir import (
     BasicQisBuilder,
+    Context,
+    FunctionType,
+    IntType,
+    PointerType,
     SimpleModule,
     Type,
-    TypeFactory,
     Value,
-    const,
-    const_getelementptr,
 )
-import re
-import pytest
-from typing import Any, Callable, List, Union
 
 
 def test_call_no_params() -> None:
     mod = SimpleModule("test", 0, 0)
-    types = mod.types
-    f = mod.add_external_function("test_function", types.function(types.void, []))
+    f = mod.add_external_function(
+        "test_function", FunctionType(Type.void(mod.context), [])
+    )
     mod.builder.call(f, [])
     assert "call void @test_function()" in mod.ir()
 
 
 def test_call_single_qubit() -> None:
     mod = SimpleModule("test", 1, 0)
-    types = mod.types
     f = mod.add_external_function(
-        "test_function", types.function(types.void, [types.qubit])
+        "test_function",
+        FunctionType(Type.void(mod.context), [pyqir.qubit_type(mod.context)]),
     )
     mod.builder.call(f, [mod.qubits[0]])
     assert "call void @test_function(%Qubit* null)" in mod.ir()
@@ -35,10 +40,10 @@ def test_call_single_qubit() -> None:
 
 def test_call_two_qubits() -> None:
     mod = SimpleModule("test", 2, 0)
-    types = mod.types
+    qubit = pyqir.qubit_type(mod.context)
     f = mod.add_external_function(
         "test_function",
-        types.function(types.void, [types.qubit, types.qubit]),
+        FunctionType(Type.void(mod.context), [qubit, qubit]),
     )
     mod.builder.call(f, [mod.qubits[0], mod.qubits[1]])
     assert (
@@ -48,60 +53,58 @@ def test_call_two_qubits() -> None:
 
 
 @pytest.mark.parametrize(
-    "get_value", [lambda types: const(types.double, 23.25), lambda _: 23.25]
+    "get_value",
+    [lambda context: pyqir.const(Type.double(context), 23.25), lambda _: 23.25],
 )
-def test_call_double(get_value: Callable[[TypeFactory], Union[Value, float]]) -> None:
+def test_call_double(get_value: Callable[[Context], Union[Value, float]]) -> None:
     mod = SimpleModule("test", 0, 0)
-    types = mod.types
     f = mod.add_external_function(
         "test_function",
-        types.function(types.void, [types.double]),
+        FunctionType(Type.void(mod.context), [Type.double(mod.context)]),
     )
-    mod.builder.call(f, [get_value(types)])
+    mod.builder.call(f, [get_value(mod.context)])
     assert "call void @test_function(double 2.325000e+01)" in mod.ir()
 
 
 @pytest.mark.parametrize(
-    "get_value", [lambda types: const(types.int(64), 42), lambda _: 42]
+    "get_value",
+    [lambda context: pyqir.const(IntType(context, 64), 42), lambda _: 42],
 )
-def test_call_int(get_value: Callable[[TypeFactory], Union[Value, int]]) -> None:
+def test_call_int(get_value: Callable[[Context], Union[Value, int]]) -> None:
     mod = SimpleModule("test", 0, 0)
-    types = mod.types
     f = mod.add_external_function(
         "test_function",
-        types.function(types.void, [types.int(64)]),
+        FunctionType(Type.void(mod.context), [IntType(mod.context, 64)]),
     )
-    mod.builder.call(f, [get_value(types)])
+    mod.builder.call(f, [get_value(mod.context)])
     assert "call void @test_function(i64 42)" in mod.ir()
 
 
 @pytest.mark.parametrize(
-    "get_value", [lambda types: const(types.bool, True), lambda _: True]
+    "get_value",
+    [lambda context: pyqir.const(IntType(context, 1), True), lambda _: True],
 )
-def test_call_bool_true(get_value: Callable[[TypeFactory], Union[Value, bool]]) -> None:
+def test_call_bool_true(get_value: Callable[[Context], Union[Value, bool]]) -> None:
     mod = SimpleModule("test", 0, 0)
-    types = mod.types
     f = mod.add_external_function(
         "test_function",
-        types.function(types.void, [types.bool]),
+        FunctionType(Type.void(mod.context), [IntType(mod.context, 1)]),
     )
-    mod.builder.call(f, [get_value(types)])
+    mod.builder.call(f, [get_value(mod.context)])
     assert "call void @test_function(i1 true)" in mod.ir()
 
 
 @pytest.mark.parametrize(
-    "get_value", [lambda types: const(types.bool, False), lambda _: False]
+    "get_value",
+    [lambda context: pyqir.const(IntType(context, 1), False), lambda _: False],
 )
-def test_call_bool_false(
-    get_value: Callable[[TypeFactory], Union[Value, bool]]
-) -> None:
+def test_call_bool_false(get_value: Callable[[Context], Union[Value, bool]]) -> None:
     mod = SimpleModule("test", 0, 0)
-    types = mod.types
     f = mod.add_external_function(
         "test_function",
-        types.function(types.void, [types.bool]),
+        FunctionType(Type.void(mod.context), [IntType(mod.context, 1)]),
     )
-    mod.builder.call(f, [get_value(types)])
+    mod.builder.call(f, [get_value(mod.context)])
     assert "call void @test_function(i1 false)" in mod.ir()
 
 
@@ -109,10 +112,9 @@ def test_call_single_result() -> None:
     mod = SimpleModule("test", 1, 1)
     qis = BasicQisBuilder(mod.builder)
     qis.mz(mod.qubits[0], mod.results[0])
-
-    types = mod.types
     f = mod.add_external_function(
-        "test_function", types.function(types.void, [types.result])
+        "test_function",
+        FunctionType(Type.void(mod.context), [pyqir.result_type(mod.context)]),
     )
     mod.builder.call(f, [mod.results[0]])
     assert "call void @test_function(%Result* null)" in mod.ir()
@@ -124,10 +126,9 @@ def test_call_two_results() -> None:
     qis.mz(mod.qubits[0], mod.results[0])
     qis.mz(mod.qubits[0], mod.results[1])
 
-    types = mod.types
+    result = pyqir.result_type(mod.context)
     f = mod.add_external_function(
-        "test_function",
-        types.function(types.void, [types.result, types.result]),
+        "test_function", FunctionType(Type.void(mod.context), [result, result])
     )
     mod.builder.call(f, [mod.results[1], mod.results[0]])
 
@@ -139,21 +140,20 @@ def test_call_two_results() -> None:
 
 def test_call_numbers() -> None:
     mod = SimpleModule("test", 0, 0)
-    types = mod.types
+    void = Type.void(mod.context)
+    i1 = IntType(mod.context, 1)
+    i64 = IntType(mod.context, 64)
+    double = Type.double(mod.context)
 
-    b = const(types.bool, True)
+    b = pyqir.const(i1, True)
     bool_rep = f"i1 true"
-    i = const(types.int(64), 42)
+    i = pyqir.const(i64, 42)
     int_rep = f"i64 42"
-    d = const(types.double, 42.42)
+    d = pyqir.const(double, 42.42)
     double_rep = "double 4.242000e+01"
 
     f = mod.add_external_function(
-        "test_function",
-        types.function(
-            types.void,
-            [types.bool, types.int(64), types.double],
-        ),
+        "test_function", FunctionType(void, [i1, i64, double])
     )
     mod.builder.call(f, [b, i, d])
     assert f"call void @test_function({bool_rep}, {int_rep}, {double_rep})" in mod.ir()
@@ -162,20 +162,19 @@ def test_call_numbers() -> None:
 @pytest.mark.parametrize(
     "get_types, args",
     [
-        (lambda types: [types.bool], ["true"]),
-        (lambda types: [types.int(64)], [1.23]),
-        (lambda types: [types.int(64)], ["123"]),
-        (lambda types: [types.double], ["1.23"]),
+        (lambda context: [IntType(context, 1)], ["true"]),
+        (lambda context: [IntType(context, 64)], [1.23]),
+        (lambda context: [IntType(context, 64)], ["123"]),
+        (lambda context: [Type.double(context)], ["1.23"]),
     ],
 )
 def test_wrong_type_conversion(
-    get_types: Callable[[TypeFactory], List[Type]], args: List[Any]
+    get_types: Callable[[Context], List[Type]], args: List[Any]
 ) -> None:
     mod = SimpleModule("test", 1, 1)
-    types = mod.types
     f = mod.add_external_function(
         "test_function",
-        types.function(types.void, get_types(types)),
+        FunctionType(Type.void(mod.context), get_types(mod.context)),
     )
     with pytest.raises(TypeError):
         mod.builder.call(f, args)
@@ -183,9 +182,9 @@ def test_wrong_type_conversion(
 
 def test_overflow_bool_value() -> None:
     mod = SimpleModule("test", 0, 0)
-    types = mod.types
-    f = mod.add_external_function("f", types.function(types.void, [types.bool]))
-    b = const(types.bool, 123)
+    i1 = IntType(mod.context, 1)
+    f = mod.add_external_function("f", FunctionType(Type.void(mod.context), [i1]))
+    b = pyqir.const(i1, 123)
     mod.builder.call(f, [b])
     assert "call void @f(i1 true)" in mod.ir()
 
@@ -193,30 +192,32 @@ def test_overflow_bool_value() -> None:
 def test_underflow_bool_value() -> None:
     mod = SimpleModule("test", 0, 0)
     with pytest.raises(OverflowError):
-        const(mod.types.bool, -123)
+        pyqir.const(IntType(mod.context, 1), -123)
 
 
 def test_overflow_bool_literal() -> None:
     mod = SimpleModule("test", 0, 0)
-    types = mod.types
-    f = mod.add_external_function("f", types.function(types.void, [types.bool]))
+    f = mod.add_external_function(
+        "f", FunctionType(Type.void(mod.context), [IntType(mod.context, 1)])
+    )
     mod.builder.call(f, [123])
     assert "call void @f(i1 true)" in mod.ir()
 
 
 def test_underflow_bool_literal() -> None:
     mod = SimpleModule("test", 0, 0)
-    types = mod.types
-    f = mod.add_external_function("f", types.function(types.void, [types.bool]))
+    f = mod.add_external_function(
+        "f", FunctionType(Type.void(mod.context), [IntType(mod.context, 1)])
+    )
     with pytest.raises(OverflowError):
         mod.builder.call(f, [-123])
 
 
 def test_overflow_int_value() -> None:
     mod = SimpleModule("test", 0, 0)
-    types = mod.types
-    f = mod.add_external_function("f", types.function(types.void, [types.int(32)]))
-    i = const(types.int(32), 2**32 + 123)
+    i32 = IntType(mod.context, 32)
+    f = mod.add_external_function("f", FunctionType(Type.void(mod.context), [i32]))
+    i = pyqir.const(i32, 2**32 + 123)
     mod.builder.call(f, [i])
     assert "call void @f(i32 123)" in mod.ir()
 
@@ -224,21 +225,23 @@ def test_overflow_int_value() -> None:
 def test_underflow_int_value() -> None:
     mod = SimpleModule("test", 0, 0)
     with pytest.raises(OverflowError):
-        const(mod.types.int(32), -(2**32) - 123)
+        pyqir.const(IntType(mod.context, 32), -(2**32) - 123)
 
 
 def test_overflow_int_literal() -> None:
     mod = SimpleModule("test", 0, 0)
-    types = mod.types
-    f = mod.add_external_function("f", types.function(types.void, [types.int(32)]))
+    f = mod.add_external_function(
+        "f", FunctionType(Type.void(mod.context), [IntType(mod.context, 32)])
+    )
     mod.builder.call(f, [2**32 + 123])
     assert "call void @f(i32 123)" in mod.ir()
 
 
 def test_underflow_int_literal() -> None:
     mod = SimpleModule("test", 0, 0)
-    types = mod.types
-    f = mod.add_external_function("f", types.function(types.void, [types.int(32)]))
+    f = mod.add_external_function(
+        "f", FunctionType(Type.void(mod.context), [IntType(mod.context, 32)])
+    )
     with pytest.raises(OverflowError):
         mod.builder.call(f, [-(2**32) - 123])
 
@@ -246,30 +249,28 @@ def test_underflow_int_literal() -> None:
 def test_64_bit_overflow() -> None:
     mod = SimpleModule("test", 0, 0)
     with pytest.raises(OverflowError):
-        const(mod.types.int(128), 2**64)
+        pyqir.const(IntType(mod.context, 128), 2**64)
 
 
 @pytest.mark.parametrize(
     "get_args",
     [
         lambda _: [],
-        lambda types: [const(types.double, 1.23)],
-        lambda types: [
-            const(types.double, 1.23),
-            const(types.bool, True),
-            const(types.bool, False),
+        lambda context: [pyqir.const(Type.double(context), 1.23)],
+        lambda context: [
+            pyqir.const(Type.double(context), 1.23),
+            pyqir.const(IntType(context, 1), True),
+            pyqir.const(IntType(context, 1), False),
         ],
     ],
 )
-def test_wrong_number_of_args(get_args: Callable[[TypeFactory], List[Value]]) -> None:
+def test_wrong_number_of_args(get_args: Callable[[Context], List[Value]]) -> None:
     mod = SimpleModule("test", 0, 0)
-    types = mod.types
-    args = get_args(types)
+    args = get_args(mod.context)
 
-    param_types: List[Type] = [types.double, types.bool]
+    param_types: List[Type] = [Type.double(mod.context), IntType(mod.context, 1)]
     f = mod.add_external_function(
-        "test_function",
-        types.function(types.void, param_types),
+        "test_function", FunctionType(Type.void(mod.context), param_types)
     )
 
     message = f"Expected {len(param_types)} arguments, got {len(args)}."
@@ -279,12 +280,9 @@ def test_wrong_number_of_args(get_args: Callable[[TypeFactory], List[Value]]) ->
 
 def test_variable() -> None:
     mod = SimpleModule("test", 0, 0)
-    types = mod.types
-    foo = mod.add_external_function("foo", types.function(types.int(64), []))
-    bar = mod.add_external_function(
-        "bar",
-        types.function(types.void, [types.int(64)]),
-    )
+    i64 = IntType(mod.context, 64)
+    foo = mod.add_external_function("foo", FunctionType(i64, []))
+    bar = mod.add_external_function("bar", FunctionType(Type.void(mod.context), [i64]))
 
     x = mod.builder.call(foo, [])
     assert x is not None
@@ -297,9 +295,10 @@ def test_variable() -> None:
 
 def test_variable_wrong_external_type() -> None:
     mod = SimpleModule("test", 0, 0)
-    types = mod.types
-    foo = mod.add_external_function("foo", types.function(types.int(64), []))
-    bar = mod.add_external_function("bar", types.function(types.void, [types.qubit]))
+    foo = mod.add_external_function("foo", FunctionType(IntType(mod.context, 64), []))
+    bar = mod.add_external_function(
+        "bar", FunctionType(Type.void(mod.context), [pyqir.qubit_type(mod.context)])
+    )
 
     x = mod.builder.call(foo, [])
     assert x is not None
@@ -313,9 +312,8 @@ def test_variable_wrong_external_type() -> None:
 
 def test_variable_wrong_angle_type() -> None:
     mod = SimpleModule("test", 1, 0)
-    types = mod.types
     qis = BasicQisBuilder(mod.builder)
-    foo = mod.add_external_function("foo", types.function(types.int(64), []))
+    foo = mod.add_external_function("foo", FunctionType(IntType(mod.context, 64), []))
 
     x = mod.builder.call(foo, [])
     assert x is not None
@@ -327,14 +325,10 @@ def test_variable_wrong_angle_type() -> None:
 
 def test_two_variables() -> None:
     mod = SimpleModule("test", 0, 0)
-    types = mod.types
-    foo = mod.add_external_function("foo", types.function(types.int(64), []))
+    i64 = IntType(mod.context, 64)
+    foo = mod.add_external_function("foo", FunctionType(i64, []))
     bar = mod.add_external_function(
-        "bar",
-        types.function(
-            types.void,
-            [types.int(64), types.int(64)],
-        ),
+        "bar", FunctionType(Type.void(mod.context), [i64, i64])
     )
 
     x = mod.builder.call(foo, [])
@@ -351,9 +345,8 @@ def test_two_variables() -> None:
 
 def test_computed_rotation() -> None:
     mod = SimpleModule("test", 1, 0)
-    types = mod.types
     qis = BasicQisBuilder(mod.builder)
-    foo = mod.add_external_function("foo", types.function(types.double, []))
+    foo = mod.add_external_function("foo", FunctionType(Type.double(mod.context), []))
 
     theta = mod.builder.call(foo, [])
     assert theta is not None
@@ -369,23 +362,36 @@ def test_computed_rotation() -> None:
 
 
 def test_record_output() -> None:
-    m = SimpleModule("test", 0, 1)
-    result_record_output = m.add_external_function(
+    mod = SimpleModule("test", 0, 1)
+    result_record_output = mod.add_external_function(
         "__quantum__rt__result_record_output",
-        m.types.function(
-            m.types.void, [m.types.result, m.types.pointer_to(m.types.int(8))]
+        FunctionType(
+            Type.void(mod.context),
+            [pyqir.result_type(mod.context), PointerType(IntType(mod.context, 8))],
         ),
     )
-    tag = m.add_global_string(b"foo")
-    i32 = m.types.int(32)
-    m.builder.call(
+    tag = mod.add_global_string(b"foo")
+    i32 = IntType(mod.context, 32)
+    mod.builder.call(
         result_record_output,
-        [m.results[0], const_getelementptr(tag, [const(i32, 0), const(i32, 0)])],
+        [
+            mod.results[0],
+            pyqir.const_getelementptr(tag, [pyqir.const(i32, 0), pyqir.const(i32, 0)]),
+        ],
     )
 
-    ir = m.ir()
+    ir = mod.ir()
     assert r'@0 = internal constant [4 x i8] c"foo\00"' in ir
     assert (
         "call void @__quantum__rt__result_record_output(%Result* null, i8* getelementptr inbounds ([4 x i8], [4 x i8]* @0, i32 0, i32 0))"
         in ir
     )
+
+
+def test_simple_module_injected_context() -> None:
+    context = Context()
+    mod = SimpleModule("test", 0, 0, context)
+    i64 = IntType(context, 64)
+    f = mod.add_external_function("f", FunctionType(Type.void(context), [i64]))
+    mod.builder.call(f, [pyqir.const(i64, 0)])
+    assert "declare void @f(i64)" in mod.ir()
diff --git a/pyqir/tests/test_generator.py b/pyqir/tests/test_generator.py
index aa68a967..d00fbff2 100644
--- a/pyqir/tests/test_generator.py
+++ b/pyqir/tests/test_generator.py
@@ -7,9 +7,11 @@
 about the generated IR.
 """
 
-from pyqir import BasicQisBuilder, SimpleModule
 import pytest
 
+import pyqir
+from pyqir import BasicQisBuilder, Context, FunctionType, IntType, SimpleModule
+
 
 def test_bell() -> None:
     module = SimpleModule("Bell circuit", num_qubits=2, num_results=2)
@@ -103,9 +105,8 @@ def test_all_gates() -> None:
 
 def test_if() -> None:
     module = SimpleModule("If", num_qubits=1, num_results=1)
-    types = module.types
     qis = BasicQisBuilder(module.builder)
-    f = module.add_external_function("f", types.function(types.int(1), []))
+    f = module.add_external_function("f", FunctionType(IntType(module.context, 1), []))
 
     b = module.builder.call(f, [])
     assert b is not None
@@ -140,7 +141,10 @@ def test_multiple_contexts() -> None:
     with pytest.raises(
         ValueError, match=r"^Some objects come from a different context\.$"
     ):
-        m1.add_external_function("f", m1.types.function(m2.types.result, []))
+        m1.add_external_function(
+            "f",
+            FunctionType(pyqir.result_type(m1.context), [pyqir.qubit_type(m2.context)]),
+        )
 
 
 def test_ir_idempotence() -> None:
@@ -171,3 +175,17 @@ def test_ir_gate_ir() -> None:
     qis.h(m.qubits[0])
     ir2 = m.ir()
     assert "call void @__quantum__qis__h__body(%Qubit* null)" in ir2
+
+
+def test_shared_context() -> None:
+    context = Context()
+
+    m1 = SimpleModule("test", 1, 1, context)
+    qis1 = BasicQisBuilder(m1.builder)
+    qis1.mz(m1.qubits[0], m1.results[0])
+
+    m2 = SimpleModule("test", 1, 1, context)
+    qis2 = BasicQisBuilder(m2.builder)
+    qis2.mz(m2.qubits[0], m2.results[0])
+
+    assert m1.ir() == m2.ir()
diff --git a/pyqir/tests/test_if.py b/pyqir/tests/test_if.py
index ec8d712e..8aca6634 100644
--- a/pyqir/tests/test_if.py
+++ b/pyqir/tests/test_if.py
@@ -1,13 +1,23 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
 
-from abc import ABCMeta, abstractmethod
-from pyqir import BasicQisBuilder, IntPredicate, SimpleModule, const
-from pyqir.evaluator import GateLogger, GateSet, NonadaptiveEvaluator
-import pytest
 import tempfile
+from abc import ABCMeta, abstractmethod
 from typing import Any, Callable, List, Optional
 
+import pytest
+
+import pyqir
+from pyqir import (
+    BasicQisBuilder,
+    FunctionType,
+    IntPredicate,
+    IntType,
+    SimpleModule,
+    Type,
+)
+from pyqir.evaluator import GateLogger, GateSet, NonadaptiveEvaluator
+
 
 class _Brancher(metaclass=ABCMeta):
     """
@@ -72,10 +82,10 @@ def if_(
 class _BoolBrancher(_Brancher):
     def __init__(self, num_queries: int) -> None:
         self._brancher = _ResultBrancher(num_queries)
-        types = self.module.types
+        context = self.module.context
         self._read_result = self._brancher.module.add_external_function(
             "__quantum__qis__read_result__body",
-            types.function(types.bool, [types.result]),
+            FunctionType(IntType(context, 1), [pyqir.result_type(context)]),
         )
 
     @property
@@ -341,7 +351,9 @@ def test_icmp_if_true() -> None:
     brancher = _BoolBrancher(1)
     x = brancher.oracle()
     module = brancher.module
-    cond = module.builder.icmp(IntPredicate.EQ, x, const(module.types.bool, 0))
+    cond = module.builder.icmp(
+        IntPredicate.EQ, x, pyqir.const(IntType(module.context, 1), 0)
+    )
 
     qis = BasicQisBuilder(module.builder)
     brancher.if_(
@@ -359,7 +371,9 @@ def test_icmp_if_false() -> None:
     brancher = _BoolBrancher(1)
     x = brancher.oracle()
     module = brancher.module
-    cond = module.builder.icmp(IntPredicate.EQ, x, const(module.types.bool, 0))
+    cond = module.builder.icmp(
+        IntPredicate.EQ, x, pyqir.const(IntType(module.context, 1), 0)
+    )
 
     qis = BasicQisBuilder(brancher.module.builder)
     brancher.if_(
@@ -379,11 +393,11 @@ def test_arithmetic_in_branch(result: bool) -> None:
     cond = brancher.oracle()
     module = brancher.module
     qis = BasicQisBuilder(module.builder)
-    i32 = module.types.int(32)
+    i32 = IntType(module.context, 32)
 
     def true() -> None:
-        four = module.builder.add(const(i32, 2), const(i32, 2))
-        cond = module.builder.icmp(IntPredicate.EQ, four, const(i32, 4))
+        four = module.builder.add(pyqir.const(i32, 2), pyqir.const(i32, 2))
+        cond = module.builder.icmp(IntPredicate.EQ, four, pyqir.const(i32, 4))
         module.builder.if_(cond, lambda: qis.x(module.qubits[0]))
 
     brancher.if_(cond, true)
@@ -398,10 +412,10 @@ def true() -> None:
 def test_call_in_branch(result: bool) -> None:
     brancher = _BoolBrancher(1)
     module = brancher.module
-    types = module.types
+    context = module.context
     x = module.add_external_function(
         "__quantum__qis__x__body",
-        types.function(types.void, [types.qubit]),
+        FunctionType(Type.void(context), [pyqir.qubit_type(context)]),
     )
 
     cond = brancher.oracle()
diff --git a/pyqir/tests/test_int_operators.py b/pyqir/tests/test_int_operators.py
index ba4fc910..a7229a3d 100644
--- a/pyqir/tests/test_int_operators.py
+++ b/pyqir/tests/test_int_operators.py
@@ -2,10 +2,21 @@
 # Licensed under the MIT License.
 
 from functools import partial
-from pyqir import Builder, IntPredicate, SimpleModule, Value, const
-import pytest
 from typing import Callable, List, Tuple
 
+import pytest
+
+import pyqir
+from pyqir import (
+    Builder,
+    FunctionType,
+    IntPredicate,
+    IntType,
+    SimpleModule,
+    Type,
+    Value,
+)
+
 _OPERATORS: List[Tuple[str, Callable[[Builder], Callable[[Value, Value], Value]]]] = [
     ("and", lambda b: b.and_),
     ("or", lambda b: b.or_),
@@ -33,10 +44,10 @@ def test_variable_variable(
     name: str, build: Callable[[Builder], Callable[[Value, Value], Value]]
 ) -> None:
     mod = SimpleModule("test " + name, 0, 0)
-    types = mod.types
-    source = mod.add_external_function("source", types.function(types.int(64), []))
-    ty = types.bool if name.startswith("icmp") else types.int(64)
-    sink = mod.add_external_function("sink", types.function(types.void, [ty]))
+    i64 = IntType(mod.context, 64)
+    source = mod.add_external_function("source", FunctionType(i64, []))
+    ty = IntType(mod.context, 1) if name.startswith("icmp") else i64
+    sink = mod.add_external_function("sink", FunctionType(Type.void(mod.context), [ty]))
     x = mod.builder.call(source, [])
     assert x is not None
     y = mod.builder.call(source, [])
@@ -51,28 +62,27 @@ def test_constant_variable(
     name: str, build: Callable[[Builder], Callable[[Value, Value], Value]]
 ) -> None:
     mod = SimpleModule("test " + name, 0, 0)
-    types = mod.types
-    source = mod.add_external_function("source", types.function(types.int(64), []))
-    ty = types.bool if name.startswith("icmp") else types.int(64)
-    sink = mod.add_external_function("sink", types.function(types.void, [ty]))
+    i64 = IntType(mod.context, 64)
+    source = mod.add_external_function("source", FunctionType(i64, []))
+    ty = IntType(mod.context, 1) if name.startswith("icmp") else i64
+    sink = mod.add_external_function("sink", FunctionType(Type.void(mod.context), [ty]))
     x = mod.builder.call(source, [])
     assert x is not None
-    y = build(mod.builder)(const(types.int(64), 1), x)
+    y = build(mod.builder)(pyqir.const(i64, 1), x)
     mod.builder.call(sink, [y])
     assert f"%1 = {name} i64 1, %0" in mod.ir()
 
 
 def test_type_mismatch() -> None:
     mod = SimpleModule("test_type_mismatch", 0, 0)
-    types = mod.types
-    source = mod.add_external_function("source", types.function(types.int(16), []))
+    i16 = IntType(mod.context, 16)
+    source = mod.add_external_function("source", FunctionType(i16, []))
     sink = mod.add_external_function(
-        "sink",
-        types.function(types.void, [types.int(16)]),
+        "sink", FunctionType(Type.void(mod.context), [i16])
     )
     x = mod.builder.call(source, [])
     assert x is not None
-    y = mod.builder.add(x, const(types.int(18), 2))
+    y = mod.builder.add(x, pyqir.const(IntType(mod.context, 18), 2))
     mod.builder.call(sink, [y])
     with pytest.raises(OSError):
         mod.ir()
diff --git a/pyqir/tests/test_qis.py b/pyqir/tests/test_qis.py
index 97fe37fb..e5916071 100644
--- a/pyqir/tests/test_qis.py
+++ b/pyqir/tests/test_qis.py
@@ -1,10 +1,13 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
 
-from pyqir import BasicQisBuilder, SimpleModule, TypeFactory, Value, const
-import pytest
 from typing import Callable, Union
 
+import pytest
+
+import pyqir
+from pyqir import BasicQisBuilder, Context, SimpleModule, Type, Value
+
 
 @pytest.mark.parametrize(
     "name, get_gate",
@@ -73,18 +76,18 @@ def test_adjoint(
 @pytest.mark.parametrize(
     "get_value",
     [
-        lambda types: const(types.double, 1.0),
+        lambda context: pyqir.const(Type.double(context), 1.0),
         lambda _: 1.0,
     ],
 )
 def test_rotated(
     name: str,
     get_gate: Callable[[BasicQisBuilder], Callable[[Union[Value, float], Value], None]],
-    get_value: Callable[[TypeFactory], Union[Value, float]],
+    get_value: Callable[[Context], Union[Value, float]],
 ) -> None:
     mod = SimpleModule("test_rotated", 1, 0)
     qis = BasicQisBuilder(mod.builder)
-    get_gate(qis)(get_value(mod.types), mod.qubits[0])
+    get_gate(qis)(get_value(mod.context), mod.qubits[0])
     call = f"call void @__quantum__qis__{name}__body(double 1.000000e+00, %Qubit* null)"
     assert call in mod.ir()