Merge pull request #972 from informalsystems/gabriela/fix-type-quanti…

…fication

Fix type quantification

CHANGELOG.md

@@ -22,8 +22,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 - Fixed a bug where sometimes static analysis would flag a mode error where
   there isn't one (#960)
-
-- Fixed the behavior of `slice` for the case `l.slice(length(l), length(l))` (#971)
+- Fixed the behavior of `slice` for the case `l.slice(length(l), length(l))`
+  (#971)
+- Fixed a bug where definitions with nested definitions could have their
+  inferred type be too general.
 
 ### Security
 

quint/src/effects/inferrer.ts

@@ -304,7 +304,7 @@ export class EffectInferrer implements IRVisitor {
 
         const nonFreeNames = effect.params.reduce(
           (names, p) => {
-            const { effectVariables: effectVariables, entityVariables: entityVariables } = effectNames(p)
+            const { effectVariables, entityVariables } = effectNames(p)
             return {
               effectVariables: new Set([...names.effectVariables, ...effectVariables]),
               entityVariables: new Set([...names.entityVariables, ...entityVariables]),

quint/src/types/constraintGenerator.ts

@@ -27,6 +27,7 @@ import {
   QuintOpDef,
   QuintStr,
   QuintVar,
+  isAnnotatedDef,
 } from '../quintIr'
 import { QuintType, typeNames } from '../quintTypes'
 import { expressionToString, rowToString, typeToString } from '../IRprinting'
@@ -69,6 +70,10 @@ export class ConstraintGeneratorVisitor implements IRVisitor {
   // Track location descriptions for error tree traces
   private location: string = ''
 
+  // A stack of free type variables and row variables for lambda expressions.
+  // Nested lambdas add new entries to the stack, and pop them when exiting.
+  private freeNames: { typeVariables: Set<string>; rowVariables: Set<string> }[] = []
+
   getResult(): [Map<bigint, ErrorTree>, Map<bigint, TypeScheme>] {
     return [this.errors, this.types]
   }
@@ -77,9 +82,17 @@ export class ConstraintGeneratorVisitor implements IRVisitor {
     this.location = `Generating constraints for ${expressionToString(e)}`
   }
 
-  exitDef(_def: QuintDef) {
+  exitDef(def: QuintDef) {
     if (this.constraints.length > 0) {
-      this.solveConstraints()
+      this.solveConstraints().map(subs => {
+        if (!isAnnotatedDef(def)) {
+          return
+        }
+
+        checkAnnotationGenerality(subs, def.typeAnnotation).mapLeft(err =>
+          this.errors.set(def.typeAnnotation?.id ?? def.id, err)
+        )
+      })
     }
   }
 
@@ -168,10 +181,18 @@ export class ConstraintGeneratorVisitor implements IRVisitor {
   }
 
   enterLambda(expr: QuintLambda) {
+    const lastParamNames = this.currentFreeNames()
+    const paramNames = {
+      typeVariables: new Set(lastParamNames.typeVariables),
+      rowVariables: new Set(lastParamNames.rowVariables),
+    }
     expr.params.forEach(p => {
       const varName = p.name === '_' ? this.freshVarGenerator.freshVar('t') : `t_${p.name}_${p.id}`
+      paramNames.typeVariables.add(varName)
       this.addToResults(p.id, right(toScheme({ kind: 'var', name: varName })))
     })
+
+    this.freeNames.push(paramNames)
   }
 
   //    Γ ∪ {p0: t0, ..., pn: tn} ⊢ e: (te, c)    t0, ..., tn are fresh
@@ -186,19 +207,21 @@ export class ConstraintGeneratorVisitor implements IRVisitor {
       return paramTypes
         .map((ts): TypeScheme => {
           const newType: QuintType = { kind: 'oper', args: ts, res: resultType.type }
-          return { ...typeNames(newType), type: newType }
+
+          return toScheme(newType)
         })
         .mapLeft(e => {
           throw new Error(`This should be impossible: Lambda variables not found: ${e.map(errorTreeToString)}`)
         })
     })
 
     this.addToResults(e.id, result)
+    this.freeNames.pop()
   }
 
   //   Γ ⊢ e1: (t1, c1)  s = solve(c1)     s(Γ ∪ {n: t1}) ⊢ e2: (t2, c2)
   // ------------------------------------------------------------------------ (LET-OPDEF)
-  //               Γ ⊢ val n = e1 { e2 }: (t2, c1 ∧ c2)
+  //               Γ ⊢ val n = e1 { e2 }: (quantify(t2), c1 ∧ c2)
   exitLet(e: QuintLet) {
     if (this.errors.size !== 0) {
       return
@@ -215,16 +238,10 @@ export class ConstraintGeneratorVisitor implements IRVisitor {
     }
 
     this.fetchResult(e.expr.id).map(t => {
-      this.addToResults(e.id, right({ ...typeNames(t.type), type: t.type }))
+      this.addToResults(e.id, right(this.quantify(t.type)))
       if (e.typeAnnotation) {
         this.constraints.push({ kind: 'eq', types: [t.type, e.typeAnnotation], sourceId: e.id })
       }
-
-      this.solveConstraints().chain(subs =>
-        checkAnnotationGenerality(subs, e.typeAnnotation).mapLeft(err =>
-          this.errors.set(e.typeAnnotation?.id ?? e.id, err)
-        )
-      )
     })
   }
 
@@ -261,7 +278,7 @@ export class ConstraintGeneratorVisitor implements IRVisitor {
         this.types = new Map<bigint, TypeScheme>(
           [...this.types.entries()].map(([id, te]) => {
             const newType = applySubstitution(this.table, subs, te.type)
-            const scheme: TypeScheme = { ...typeNames(newType), type: newType }
+            const scheme: TypeScheme = this.quantify(newType)
             return [id, scheme]
           })
         )
@@ -310,6 +327,24 @@ export class ConstraintGeneratorVisitor implements IRVisitor {
     const subs = compose(this.table, typeSubs, rowSubs)
     return applySubstitution(this.table, subs, t.type)
   }
+
+  private currentFreeNames(): { typeVariables: Set<string>; rowVariables: Set<string> } {
+    return (
+      this.freeNames[this.freeNames.length - 1] ?? {
+        typeVariables: new Set(),
+        rowVariables: new Set(),
+      }
+    )
+  }
+
+  private quantify(type: QuintType): TypeScheme {
+    const freeNames = this.currentFreeNames()
+    const nonFreeNames = {
+      typeVariables: new Set([...typeNames(type).typeVariables].filter(name => !freeNames.typeVariables.has(name))),
+      rowVariables: new Set([...typeNames(type).rowVariables].filter(name => !freeNames.rowVariables.has(name))),
+    }
+    return { ...nonFreeNames, type }
+  }
 }
 
 function checkAnnotationGenerality(

quint/test/types/inferrer.test.ts

@@ -155,6 +155,16 @@ describe('inferTypes', () => {
     ])
   })
 
+  it('keeps track of free variables in nested scopes (#966)', () => {
+    const quintModule = buildModuleWithDefs(['def f(a) = a == "x"', 'def g(b) = val nested = (1,2) { f(b) }'])
+
+    const [errors, types] = inferTypesForModule(quintModule)
+    assert.isEmpty(errors, `Should find no errors, found: ${[...errors.values()].map(errorTreeToString)}`)
+
+    const stringTypes = Array.from(types.entries()).map(([id, type]) => [id, typeSchemeToString(type)])
+    assert.includeDeepMembers(stringTypes, [[14n, '(str) => bool']])
+  })
+
   it('considers annotations', () => {
     const quintModule = buildModuleWithDefs(['def e(p): (int) => int = p'])