Consolidate row simplification utility

quint/src/graphics.ts

@@ -32,7 +32,8 @@ import { zerog } from './idGenerator'
 import { ConcreteFixedRow, QuintType, Row } from './ir/quintTypes'
 import { TypeScheme } from './types/base'
 import { canonicalTypeScheme } from './types/printing'
-import { declarationToString, flattenRow, qualifierToString, rowToString } from './ir/IRprinting'
+import { declarationToString, qualifierToString, rowToString } from './ir/IRprinting'
+import { simplifyRow } from './types/simplification'
 
 /**
  * An abstraction of a Console of bounded text width.
@@ -205,23 +206,27 @@ export function prettyQuintType(type: QuintType): Doc {
   }
 }
 
-function prettyRow(row: Row, showFieldName = true): Doc {
-  const [fields, other] = flattenRow(row)
+function prettyRow(r: Row, showFieldName = true): Doc {
+  const row = simplifyRow(r)
+  const fields = row.kind === 'row' ? row.fields : []
+  const other = row.kind === 'row' ? row.other : undefined
 
   const fieldsDocs = fields.map(f => {
     const prefix = showFieldName ? `${f.fieldName}: ` : ''
     return group([text(prefix), prettyQuintType(f.fieldType)])
   })
-  const otherDoc = other.kind === 'var' ? [text(` | ${other.name}`)] : []
+  const otherDoc = other?.kind === 'var' ? [text(` | ${other.name}`)] : []
 
   return group([nest('  ', [linebreak, docJoin([text(','), line()], fieldsDocs)]), ...otherDoc, linebreak])
 }
 
-function prettySumRow(row: ConcreteFixedRow): Doc {
-  const [fields, other] = flattenRow(row)
+function prettySumRow(r: ConcreteFixedRow): Doc {
+  const row = simplifyRow(r)
+  const fields = row.kind === 'row' ? row.fields : []
+  const other = row.kind === 'row' ? row.other : undefined
 
   const fieldsDocs = fields.map(f => {
-    if (other.kind === 'empty') {
+    if (other?.kind === 'empty') {
       return group(text(f.fieldName))
     } else {
       return group([text(`${f.fieldName}(`), prettyQuintType(f.fieldType), text(')')])
@@ -230,6 +235,7 @@ function prettySumRow(row: ConcreteFixedRow): Doc {
 
   return group([nest('  ', [linebreak, docJoin([text(','), line()], fieldsDocs)]), linebreak])
 }
+
 /**
  * Print an execution frame and its children recursively.
  * Since this function is printing a tree, we need precise text alignment.

quint/src/ir/IRTransformer.ts

@@ -504,7 +504,7 @@ function transformExpression(transformer: IRTransformer, expr: ir.QuintEx): ir.Q
  *
  * @returns the transformed Quint row
  */
-function transformRow(transformer: IRTransformer, row: t.Row): t.Row {
+export function transformRow(transformer: IRTransformer, row: t.Row): t.Row {
   let newRow = row
   if (transformer.enterRow) {
     newRow = transformer.enterRow(newRow)

quint/src/ir/IRprinting.ts

@@ -13,7 +13,7 @@
  */
 
 import { OpQualifier, QuintDeclaration, QuintDef, QuintEx, QuintModule, isAnnotatedDef } from './quintIr'
-import { EmptyRow, QuintSumType, QuintType, Row, RowField, VarRow, isTheUnit } from './quintTypes'
+import { QuintSumType, QuintType, Row, RowField, isTheUnit } from './quintTypes'
 import { TypeScheme } from '../types/base'
 import { typeSchemeToString } from '../types/printing'
 
@@ -269,16 +269,3 @@ function rowFieldsToString(r: Row, showFieldName = true): string {
     }
   }
 }
-
-export function flattenRow(r: Row): [{ fieldName: string; fieldType: QuintType }[], VarRow | EmptyRow] {
-  switch (r.kind) {
-    case 'empty':
-      return [[], r]
-    case 'var':
-      return [[], r]
-    case 'row': {
-      const [fields, other] = flattenRow(r.other)
-      return [[...r.fields, ...fields], other]
-    }
-  }
-}

quint/src/types/constraintSolver.ts

@@ -20,6 +20,7 @@ import { Constraint } from './base'
 import { Substitutions, applySubstitution, applySubstitutionToConstraint, compose } from './substitutions'
 import { unzip } from 'lodash'
 import { LookupTable } from '../names/base'
+import { simplifyRow } from './simplification'
 
 /*
  * Try to solve a constraint by unifying all pairs of types in equality
@@ -251,24 +252,6 @@ function chainUnifications(table: LookupTable, types1: QuintType[], types2: Quin
   }, right([]))
 }
 
-function simplifyRow(r: Row): Row {
-  if (r.kind !== 'row') {
-    return r
-  }
-
-  let result = r
-  const otherSimplified = simplifyRow(r.other)
-  if (otherSimplified.kind === 'row') {
-    result = { kind: 'row', fields: r.fields.concat(otherSimplified.fields), other: otherSimplified.other }
-  }
-
-  if (result.fields.length > 0) {
-    return result
-  } else {
-    return otherSimplified
-  }
-}
-
 function tryToUnpack(
   location: string,
   types1: QuintType[],

quint/src/types/inferrer.ts

@@ -20,6 +20,7 @@ import { QuintDeclaration } from '../ir/quintIr'
 import { TypeScheme } from './base'
 import { ConstraintGeneratorVisitor } from './constraintGenerator'
 import { solveConstraint } from './constraintSolver'
+import { simplify } from './simplification'
 
 export type TypeInferenceResult = [Map<bigint, ErrorTree>, Map<bigint, TypeScheme>]
 
@@ -40,6 +41,7 @@ export class TypeInferrer extends ConstraintGeneratorVisitor {
     declarations.forEach(decl => {
       walkDeclaration(this, decl)
     })
-    return [this.errors, this.types]
+    const simplifiedTypes = new Map([...this.types.entries()].map(([id, t]) => [id, simplify(t)]))
+    return [this.errors, simplifiedTypes]
   }
 }

quint/src/types/simplification.ts

@@ -0,0 +1,65 @@
+/* ----------------------------------------------------------------------------------
+ * Copyright (c) Informal Systems 2023. All rights reserved.
+ * Licensed under the Apache 2.0.
+ * See License.txt in the project root for license information.
+ * --------------------------------------------------------------------------------- */
+
+/**
+ * Simplification for quint types
+ *
+ * @author Gabriela Moreira
+ *
+ * @module
+ */
+
+import { IRTransformer, transformRow, transformType } from '../ir/IRTransformer'
+import { Row } from '../ir/quintTypes'
+import { TypeScheme } from './base'
+
+/**
+ * Simplifies a type scheme by flattening all type rows.
+ *
+ * A row like `{ a: int | { b: bool | r } }` becomes
+ * `{ a: int, b: bool | r }`.
+ *
+ * @param typeScheme - The type scheme to be simplified
+ *
+ * @returns The simplified type scheme
+ */
+export function simplify(typeScheme: TypeScheme): TypeScheme {
+  const simplifier = new RowSimplifier()
+  return { ...typeScheme, type: transformType(simplifier, typeScheme.type) }
+}
+
+/**
+ * Simplifies a row type.
+ *
+ * A row like `{ a: int | { b: bool | r } }` becomes
+ * `{ a: int, b: bool | r }`.
+ *
+ * @param row - The tow to be simplified
+ *
+ * @returns The simplified row
+ */
+export function simplifyRow(row: Row): Row {
+  const simplifier = new RowSimplifier()
+  return transformRow(simplifier, row)
+}
+
+class RowSimplifier implements IRTransformer {
+  exitRow(row: Row): Row {
+    if (row.kind !== 'row') {
+      return row
+    }
+
+    if (row.fields.length === 0) {
+      return row.other
+    }
+
+    if (row.other.kind !== 'row') {
+      return row
+    }
+
+    return { ...row, fields: row.fields.concat(row.other.fields), other: row.other.other }
+  }
+}

quint/test/types/simplification.test.ts

@@ -0,0 +1,56 @@
+import { describe, it } from 'mocha'
+import { assert } from 'chai'
+import { typeSchemeToString } from '../../src/types/printing'
+import { simplify } from '../../src/types/simplification'
+import { toScheme } from '../../src/types/base'
+import { QuintRecordType } from '../../src'
+
+describe('simplify', () => {
+  it('flattens nested row tails', () => {
+    // { a: int | { b: bool | { c: str | v } } }
+    const type: QuintRecordType = {
+      kind: 'rec',
+      fields: {
+        kind: 'row',
+        fields: [{ fieldName: 'a', fieldType: { kind: 'int' } }],
+        other: {
+          kind: 'row',
+          fields: [{ fieldName: 'b', fieldType: { kind: 'bool' } }],
+          other: {
+            kind: 'row',
+            fields: [{ fieldName: 'c', fieldType: { kind: 'str' } }],
+            other: { kind: 'var', name: 'v' },
+          },
+        },
+      },
+    }
+
+    const simplifiedType = simplify(toScheme(type))
+
+    assert.deepEqual(typeSchemeToString(simplifiedType), '{ a: int, b: bool, c: str | v }')
+  })
+
+  it('simplifies rows with empty fields', () => {
+    // { | { b: bool | { c: str | v } } }
+    const type: QuintRecordType = {
+      kind: 'rec',
+      fields: {
+        kind: 'row',
+        fields: [],
+        other: {
+          kind: 'row',
+          fields: [{ fieldName: 'b', fieldType: { kind: 'bool' } }],
+          other: {
+            kind: 'row',
+            fields: [{ fieldName: 'c', fieldType: { kind: 'str' } }],
+            other: { kind: 'var', name: 'v' },
+          },
+        },
+      },
+    }
+
+    const simplifiedType = simplify(toScheme(type))
+
+    assert.deepEqual(typeSchemeToString(simplifiedType), '{ b: bool, c: str | v }')
+  })
+})