Fix out-of-memory in typeck on nested Ty::Arrow (#2336)

This change updates the `Ty::Arrow` enum variant to use `Rc` which
allows sharing of instances of callable types within the type
infrastructure. To allow mutability during type inference, the
underlying `Arrow` struct now uses `RefCell` for it's contained type and
relies on runtime borrow checking for mutability.

Fixes #2302

---------

Co-authored-by: Ian Davis <[email protected]>

Author: swernli

compiler/qsc/Cargo.toml

@@ -80,3 +80,7 @@ harness = false
 [[bench]]
 name = "rca"
 harness = false
+
+[[bench]]
+name = "typeck"
+harness = false

compiler/qsc/benches/typeck.rs

@@ -0,0 +1,43 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+use criterion::{criterion_group, criterion_main, Criterion};
+use indoc::indoc;
+use qsc::{
+    interpret::Interpreter, LanguageFeatures, PackageType, SourceMap, TargetCapabilityFlags,
+};
+
+pub fn deep_nested_callable_generics(c: &mut Criterion) {
+    c.bench_function("Deeply nested callable generics", |b| {
+        b.iter(move || {
+            let sources = SourceMap::new(
+                [("none".into(), "".into())],
+                Some(
+                    indoc! {"{
+                    function swap<'T1, 'T2>(op : ('T1, 'T2) -> ('T2, 'T1), input: ('T1, 'T2)) : ('T2, 'T1) {
+                        op(input)
+                    }
+                    function nested_swaps() : Unit {
+                        let a = swap(swap(swap(swap(swap(_)))));
+                    }
+            }"}
+                    .into(),
+                ),
+            );
+            let (std_id, store) = qsc::compile::package_store_with_stdlib(TargetCapabilityFlags::all());
+
+            assert!(Interpreter::new(
+                sources,
+                PackageType::Exe,
+                TargetCapabilityFlags::all(),
+                LanguageFeatures::default(),
+                store,
+                &[(std_id, None)],
+            )
+            .is_err(), "code should fail with type error");
+        });
+    });
+}
+
+criterion_group!(benches, deep_nested_callable_generics);
+criterion_main!(benches);

compiler/qsc_frontend/src/closure.rs

@@ -13,7 +13,7 @@ use qsc_hir::{
     visit::{self, Visitor},
 };
 use rustc_hash::{FxHashMap, FxHashSet};
-use std::iter;
+use std::{iter, rc::Rc};
 
 pub(super) struct Lambda {
     pub(super) kind: CallableKind,
@@ -162,27 +162,28 @@ pub(super) fn partial_app_block(
     callee: Expr,
     arg: Expr,
     app: PartialApp,
-    arrow: Arrow,
+    arrow: Rc<Arrow>,
     span: Span,
 ) -> Block {
     let call = Expr {
         id: NodeId::default(),
         span,
-        ty: (*arrow.output).clone(),
+        ty: arrow.output.borrow().clone(),
         kind: ExprKind::Call(Box::new(callee), Box::new(arg)),
     };
     let lambda = Lambda {
         kind: arrow.kind,
         functors: arrow
             .functors
+            .borrow()
             .expect_value("lambda type should have concrete functors"),
         input: app.input,
         body: call,
     };
     let closure = Expr {
         id: NodeId::default(),
         span,
-        ty: Ty::Arrow(Box::new(arrow.clone())),
+        ty: Ty::Arrow(arrow.clone()),
         kind: close(lambda),
     };
 
@@ -195,7 +196,7 @@ pub(super) fn partial_app_block(
     Block {
         id: NodeId::default(),
         span,
-        ty: Ty::Arrow(Box::new(arrow)),
+        ty: Ty::Arrow(arrow),
         stmts,
     }
 }

compiler/qsc_frontend/src/compile/tests.rs

@@ -3,7 +3,7 @@
 
 mod multiple_packages;
 
-use std::sync::Arc;
+use std::{rc::Rc, sync::Arc};
 
 use super::{compile, longest_common_prefix, CompileUnit, Error, PackageStore, SourceMap};
 use crate::compile::TargetCapabilityFlags;
@@ -356,7 +356,7 @@ fn insert_core_call() {
             let callee = Expr {
                 id: NodeId::default(),
                 span: Span::default(),
-                ty: Ty::Arrow(Box::new(allocate_ty)),
+                ty: Ty::Arrow(Rc::new(allocate_ty)),
                 kind: ExprKind::Var(Res::Item(allocate.id), Vec::new()),
             };
 

compiler/qsc_frontend/src/lower.rs

@@ -739,7 +739,7 @@ impl With<'_> {
             ast::ExprKind::Block(block) => hir::ExprKind::Block(self.lower_block(block)),
             ast::ExprKind::Call(callee, arg) => match &ty {
                 Ty::Arrow(arrow) if is_partial_app(arg) => hir::ExprKind::Block(
-                    self.lower_partial_app(callee, arg, (**arrow).clone(), expr.span),
+                    self.lower_partial_app(callee, arg, arrow.clone(), expr.span),
                 ),
                 _ => hir::ExprKind::Call(
                     Box::new(self.lower_expr(callee)),
@@ -779,6 +779,7 @@ impl With<'_> {
                 let functors = if let Ty::Arrow(arrow) = &ty {
                     arrow
                         .functors
+                        .borrow()
                         .expect_value("lambda type should have concrete functors")
                 } else {
                     FunctorSetValue::Empty
@@ -881,7 +882,7 @@ impl With<'_> {
         &mut self,
         callee: &ast::Expr,
         arg: &ast::Expr,
-        arrow: Arrow,
+        arrow: Rc<Arrow>,
         span: Span,
     ) -> hir::Block {
         let callee = self.lower_expr(callee);

compiler/qsc_frontend/src/lower/tests.rs

@@ -1356,7 +1356,7 @@ fn lambda_operation_empty_closure() {
                                     Pat 7 [69-70] [Type Qubit]: Bind: Ident 8 [69-70] "q"
                                     QubitInit 9 [73-80] [Type Qubit]: Single
                                 Stmt 10 [90-95]: Expr: Expr 11 [90-95] [Type Unit]: Call:
-                                    Expr 12 [90-92] [Type (Qubit => Unit)]: Var: Local 3
+                                    Expr 12 [90-92] [Type (Qubit => Unit is Param<0>)]: Var: Local 3
                                     Expr 13 [93-94] [Type Qubit]: Var: Local 8
                         adj: <none>
                         ctl: <none>
@@ -1436,7 +1436,7 @@ fn lambda_operation_closure() {
                         body: SpecDecl 9 [80-130]: Impl:
                             Block 10 [122-130] [Type Result]:
                                 Stmt 11 [124-128]: Expr: Expr 12 [124-128] [Type Result]: Call:
-                                    Expr 13 [124-126] [Type (Unit => Result)]: Var: Local 8
+                                    Expr 13 [124-126] [Type (Unit => Result is Param<0>)]: Var: Local 8
                                     Expr 14 [126-128] [Type Unit]: Unit
                         adj: <none>
                         ctl: <none>

compiler/qsc_frontend/src/typeck/convert.rs

@@ -2,7 +2,7 @@
 // Licensed under the MIT License.
 
 //! Ascribe types to the AST and output HIR items. Put another way, converts the AST to the HIR.
-use std::rc::Rc;
+use std::{cell::RefCell, rc::Rc};
 
 use crate::resolve::{self, Names};
 
@@ -61,11 +61,11 @@ pub(crate) fn ty_from_ast(
             let functors = functors
                 .as_ref()
                 .map_or(FunctorSetValue::Empty, |f| eval_functor_expr(f.as_ref()));
-            let ty = Ty::Arrow(Box::new(Arrow {
+            let ty = Ty::Arrow(Rc::new(Arrow {
                 kind: callable_kind_from_ast(*kind),
-                input: Box::new(input),
-                output: Box::new(output),
-                functors: FunctorSet::Value(functors),
+                input: RefCell::new(input),
+                output: RefCell::new(output),
+                functors: RefCell::new(FunctorSet::Value(functors)),
             }));
             (ty, errors)
         }
@@ -157,9 +157,9 @@ pub(super) fn ast_ty_def_cons(
     let (input, errors) = ast_ty_def_base(names, def);
     let ty = Arrow {
         kind: hir::CallableKind::Function,
-        input: Box::new(input),
-        output: Box::new(Ty::Udt(ty_name.clone(), hir::Res::Item(id))),
-        functors: FunctorSet::Value(FunctorSetValue::Empty),
+        input: RefCell::new(input),
+        output: RefCell::new(Ty::Udt(ty_name.clone(), hir::Res::Item(id))),
+        functors: RefCell::new(FunctorSet::Value(FunctorSetValue::Empty)),
     };
     let scheme = Scheme::new(Vec::new(), Box::new(ty));
     (scheme, errors)
@@ -274,9 +274,9 @@ pub(super) fn scheme_for_ast_callable(
 
     let ty = Arrow {
         kind,
-        input: Box::new(input),
-        output: Box::new(output),
-        functors: FunctorSet::Value(ast_callable_functors(callable)),
+        input: RefCell::new(input),
+        output: RefCell::new(output),
+        functors: RefCell::new(FunctorSet::Value(ast_callable_functors(callable))),
     };
 
     (Scheme::new(type_parameters, Box::new(ty)), errors)
@@ -290,15 +290,24 @@ pub(crate) fn synthesize_functor_params(
 ) -> Vec<HirTypeParameter> {
     match ty {
         Ty::Array(item) => synthesize_functor_params(next_param, item),
-        Ty::Arrow(arrow) => match arrow.functors {
-            FunctorSet::Value(functors) if arrow.kind == hir::CallableKind::Operation => {
-                let param = HirTypeParameter::Functor(functors);
-                arrow.functors = FunctorSet::Param(*next_param, functors);
-                *next_param = next_param.successor();
-                vec![param]
+        Ty::Arrow(arrow) => {
+            let functors = *arrow.functors.borrow();
+            match functors {
+                FunctorSet::Value(functors) if arrow.kind == hir::CallableKind::Operation => {
+                    let param = HirTypeParameter::Functor(functors);
+                    // This uses `RefCell::replace` to update the functor set within the reference counted arrow.
+                    // This should be safe as no other references to the functors should be held, but if they are
+                    // this will panic at runtime.
+                    // Note the `next_param` local binding is needed to ensure that the `functors` are not borrowed
+                    // while we are trying to replace them.
+                    let new_functors = FunctorSet::Param(*next_param, functors);
+                    arrow.functors.replace(new_functors);
+                    *next_param = next_param.successor();
+                    vec![param]
+                }
+                _ => Vec::new(),
             }
-            _ => Vec::new(),
-        },
+        }
         Ty::Tuple(items) => items
             .iter_mut()
             .flat_map(|item| synthesize_functor_params(next_param, item))