feat(ir): Introduce branded types

Closes #191

Author: jubnzv

CHANGELOG.md

@@ -21,6 +21,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Changed
 - Improved and optimized the test suite: PR [#184](https://github.com/nowarp/misti/pull/184)
+- Introduced the branded type pattern to improve type safety: Issue [#191](https://github.com/nowarp/misti/issues/191)
 
 ## [0.4.2] - 2024-10-12
 

examples/implicit-init/implicitInit.ts

@@ -1,5 +1,5 @@
 import { ASTDetector } from "../../src/detectors/detector";
-import { CompilationUnit } from "../../src/internals/ir";
+import { CompilationUnit, FunctionName } from "../../src/internals/ir";
 import { MistiTactWarning, Severity } from "../../src/internals/warnings";
 
 /**
@@ -32,7 +32,7 @@ export class ImplicitInit extends ASTDetector {
 
   async check(cu: CompilationUnit): Promise<MistiTactWarning[]> {
     return Array.from(cu.contracts).reduce((foundErrors, [_, contract]) => {
-      if (!cu.findMethodCFGByName(contract.name, "init")) {
+      if (!cu.findMethodCFGByName(contract.name, "init" as FunctionName)) {
         const err = this.makeWarning(
           `Contract ${contract.name} doesn't define an init function`,
           contract.ref,

src/detectors/builtin/argCopyMutation.ts

@@ -64,7 +64,7 @@ export class ArgCopyMutation extends ASTDetector {
             fun.kind === "function_def" &&
             this.usedInAllReturns(
               argName,
-              returnStatements.get(idText(fun.name))!,
+              returnStatements.get(idText(fun.name) as FunctionName)!,
             )
           ) {
             return;

src/internals/ir/astStore.ts

@@ -1,5 +1,6 @@
 import { InternalException } from "../exceptions";
 import { mergeMaps } from "../util";
+import { FunctionName } from "./types";
 import {
   AstAsmFunctionDef,
   AstConstantDef,
@@ -31,8 +32,6 @@ export type AstItemParams = Partial<{
   filename?: string;
 }>;
 
-type FunctionName = string;
-type MethodName = string;
 type Filename = string;
 
 /**
@@ -536,12 +535,14 @@ export class TactASTStore {
    */
   public getReturnTypes(): Map<FunctionName, AstType | null> {
     const result = new Map<FunctionName, AstType | null>();
-    this.asmFunctions.forEach((f) => result.set(idText(f.name), f.return));
+    this.asmFunctions.forEach((f) =>
+      result.set(idText(f.name) as FunctionName, f.return),
+    );
     this.functions.forEach((f) => {
       if (!this.isContractItem(f.id))
         result.set(
           // Other kinds of functions cannot be present on the top-level
-          idText((f as AstFunctionDef).name),
+          idText((f as AstFunctionDef).name) as FunctionName,
           (f as AstFunctionDef).return,
         );
     });
@@ -557,8 +558,8 @@ export class TactASTStore {
     entryId: AstNode["id"],
     withTraits: boolean = true,
     visited = new Set<number>(),
-  ): Map<MethodName, AstType | null> {
-    let result = new Map<MethodName, AstType | null>();
+  ): Map<FunctionName, AstType | null> {
+    let result = new Map<FunctionName, AstType | null>();
 
     // Avoid recursion if used on AST before typechecking
     if (visited.has(entryId)) return result;
@@ -570,7 +571,7 @@ export class TactASTStore {
       this.functions.forEach((f) => {
         // Don't consider receivers/inits since we cannot call them from the contract
         if (f.kind === "function_def" && contractEntries.has(f.id))
-          result.set(idText(f.name), f.return);
+          result.set(idText(f.name) as FunctionName, f.return);
       });
     }
 

src/internals/ir/builders/ir.ts

@@ -1,12 +1,12 @@
 import {
+  ContractIdx,
   BasicBlock,
   BasicBlockIdx,
   BasicBlockKind,
   CFG,
   CFGIdx,
   CompilationUnit,
   Contract,
-  ContractIdx,
   ContractName,
   Edge,
   FunctionKind,
@@ -113,14 +113,15 @@ export class TactIRBuilder {
   private registerFunctions(): void {
     this.functionIndexes = this.ast.functions.reduce((acc, fun) => {
       if (fun.kind == "function_def") {
+        const funName = fun.name.text as FunctionName;
         const idx = this.registerCFGIdx(
-          fun.name.text,
+          funName,
           fun.id,
           "function",
           fun.loc.origin,
           fun.loc,
         );
-        acc.set(fun.name.text, idx);
+        acc.set(funName, idx);
       }
       return acc;
     }, new Map<FunctionName, CFGIdx>());
@@ -133,12 +134,13 @@ export class TactIRBuilder {
   private createFunctions(): Map<CFGIdx, CFG> {
     return this.ast.functions.reduce((acc, fun) => {
       if (fun.kind == "function_def") {
-        const idx = this.functionIndexes.get(fun.name.text)!;
+        const funName = fun.name.text as FunctionName;
+        const idx = this.functionIndexes.get(funName)!;
         acc.set(
           idx,
           this.createCFGFromStatements(
             idx,
-            fun.name.text,
+            funName,
             fun.id,
             "function",
             fun.loc.origin,
@@ -157,13 +159,21 @@ export class TactIRBuilder {
   private getMethodInfo(
     decl: AstContractDeclaration,
     contractId: number,
-  ): [string | undefined, FunctionKind | undefined, AstStatement[] | null] {
+  ): [
+    FunctionName | undefined,
+    FunctionKind | undefined,
+    AstStatement[] | null,
+  ] {
     return decl.kind === "function_def"
-      ? [decl.name.text, "method", decl.statements]
+      ? [decl.name.text as FunctionName, "method", decl.statements]
       : decl.kind === "contract_init"
-        ? [`init_${contractId}`, "method", decl.statements]
+        ? [`init_${contractId}` as FunctionName, "method", decl.statements]
         : decl.kind === "receiver"
-          ? [generateReceiveName(decl), "receive", decl.statements]
+          ? [
+              generateReceiveName(decl) as FunctionName,
+              "receive",
+              decl.statements,
+            ]
           : [undefined, undefined, null];
   }
 
@@ -173,7 +183,7 @@ export class TactIRBuilder {
   private registerContracts(): void {
     this.methodIndexes = this.ast.types.reduce((acc, entry) => {
       if (entry.kind == "contract") {
-        const contractName = entry.name.text;
+        const contractName = entry.name.text as ContractName;
         const methodsMap = entry.declarations.reduce((methodAcc, decl) => {
           const [name, kind, _] = this.getMethodInfo(decl, entry.id);
           // NOTE: We don't create CFG entries for asm functions.
@@ -201,7 +211,7 @@ export class TactIRBuilder {
   private createContracts(): Map<ContractIdx, Contract> {
     return this.ast.types.reduce((acc, entry) => {
       if (entry.kind == "contract") {
-        const contractName = entry.name.text;
+        const contractName = entry.name.text as ContractName;
         const methodsMap = this.methodIndexes.get(contractName)!;
         const methodCFGs = entry.declarations.reduce((methodAcc, decl) => {
           const [name, kind, stmts] = this.getMethodInfo(decl, entry.id);
@@ -281,11 +291,15 @@ export class TactIRBuilder {
     parentCalls: Set<CFGIdx> = new Set(),
   ): Set<CFGIdx> {
     switch (expr.kind) {
-      case "method_call": // method
+      case "method_call":
         if (expr.self.kind === "id" && isSelfId(expr.self)) {
-          const contractMethods = this.methodIndexes.get(expr.self.text);
+          const contractMethods = this.methodIndexes.get(
+            expr.self.text as ContractName,
+          );
           if (contractMethods) {
-            const methodIdx = contractMethods.get(expr.method.text);
+            const methodIdx = contractMethods.get(
+              expr.method.text as FunctionName,
+            );
             if (methodIdx !== undefined) {
               parentCalls.add(methodIdx);
             } else {
@@ -310,8 +324,10 @@ export class TactIRBuilder {
         expr.args.forEach((arg) => this.collectFunctionCalls(arg, parentCalls));
         this.collectFunctionCalls(expr.self, parentCalls);
         break;
-      case "static_call": // free function
-        const funcIdx = this.functionIndexes.get(expr.function.text);
+      case "static_call":
+        const funcIdx = this.functionIndexes.get(
+          expr.function.text as FunctionName,
+        );
         if (funcIdx !== undefined) {
           parentCalls.add(funcIdx);
         }

src/internals/ir/cfg.ts

@@ -6,11 +6,15 @@
  */
 import { TactASTStore } from "./astStore";
 import { IdxGenerator } from "./indices";
-import { BasicBlockIdx, CFGIdx, EdgeIdx, FunctionName } from "./types";
+import { FunctionName } from "./types";
 import { InternalException } from "../exceptions";
 import { AstStatement, SrcInfo } from "@tact-lang/compiler/dist/grammar/ast";
 import { ItemOrigin } from "@tact-lang/compiler/dist/grammar/grammar";
 
+export type EdgeIdx = number & { readonly __brand: unique symbol };
+export type BasicBlockIdx = number & { readonly __brand: unique symbol };
+export type CFGIdx = number & { readonly __brand: unique symbol };
+
 /**
  * Represents an edge in a Control Flow Graph (CFG), connecting two basic blocks.
  * Each edge signifies a potential flow of control from one statement to another.
@@ -24,7 +28,7 @@ export class Edge {
     public src: BasicBlockIdx,
     public dst: BasicBlockIdx,
   ) {
-    this.idx = IdxGenerator.next("cfg_edge");
+    this.idx = IdxGenerator.next("cfg_edge") as EdgeIdx;
   }
 }
 
@@ -61,12 +65,12 @@ export type BasicBlockKind =
 export class BasicBlock {
   public idx: BasicBlockIdx;
   constructor(
-    public stmtID: number,
+    public stmtID: AstStatement["id"],
     public kind: BasicBlockKind,
     public srcEdges: Set<EdgeIdx> = new Set<EdgeIdx>(),
     public dstEdges: Set<EdgeIdx> = new Set<EdgeIdx>(),
   ) {
-    this.idx = IdxGenerator.next("cfg_bb");
+    this.idx = IdxGenerator.next("cfg_bb") as BasicBlockIdx;
   }
 
   /**
@@ -122,7 +126,7 @@ export class CFG {
     public ref: SrcInfo,
     idx: CFGIdx | undefined = undefined,
   ) {
-    this.idx = idx ? idx : IdxGenerator.next("cfg");
+    this.idx = idx ? idx : (IdxGenerator.next("cfg") as CFGIdx);
     this.bbsMap = new Map();
     this.initializeMapping(this.bbsMap, nodes);
     this.edgesMap = new Map();
@@ -163,7 +167,7 @@ export class CFG {
     isSrc: boolean,
   ): BasicBlock[] | undefined {
     return Array.from(edgeIdxs).reduce(
-      (acc, srcIdx: BasicBlockIdx) => {
+      (acc, srcIdx) => {
         if (acc === undefined) {
           return undefined;
         }

src/internals/ir/imports.ts

@@ -3,8 +3,8 @@ import { SrcInfo } from "@tact-lang/compiler/dist/grammar/ast";
 import { ItemOrigin } from "@tact-lang/compiler/dist/grammar/grammar";
 import path from "path";
 
-export type ImportNodeIdx = number;
-export type ImportEdgeIdx = number;
+export type ImportNodeIdx = number & { readonly __brand: unique symbol };
+export type ImportEdgeIdx = number & { readonly __brand: unique symbol };
 export type ImportLanguage = "tact" | "func";
 export type ImportDirection = "forward" | "backward";
 
@@ -27,7 +27,7 @@ export class ImportNode {
     public inEdges: Set<ImportEdgeIdx> = new Set(),
     public outEdges: Set<ImportEdgeIdx> = new Set(),
   ) {
-    this.idx = IdxGenerator.next("import_node");
+    this.idx = IdxGenerator.next("import_node") as ImportNodeIdx;
   }
 }
 
@@ -42,7 +42,7 @@ export class ImportEdge {
     /** Source location of the `import` statement. */
     public loc: SrcInfo,
   ) {
-    this.idx = IdxGenerator.next("import_edge");
+    this.idx = IdxGenerator.next("import_edge") as ImportEdgeIdx;
   }
 }
 

src/internals/ir/index.ts

@@ -1,6 +1,6 @@
 export * from "./cfg";
 export * from "./astStore";
 export * from "./imports";
-export { CompilationUnit, Contract } from "./ir";
+export * from "./ir";
 export * from "./types";
 export * from "./builders/";

src/internals/ir/ir.ts

@@ -3,20 +3,15 @@
  *
  * @packageDocumentation
  */
+import { CFGIdx, ContractName, FunctionName, ProjectName } from ".";
 import { TactASTStore } from "./astStore";
 import { BasicBlock, CFG } from "./cfg";
 import { ImportGraph } from "./imports";
 import { IdxGenerator } from "./indices";
-import {
-  BasicBlockIdx,
-  CFGIdx,
-  ContractIdx,
-  ContractName,
-  FunctionName,
-  ProjectName,
-} from "./types";
 import { AstStatement, SrcInfo } from "@tact-lang/compiler/dist/grammar/ast";
 
+export type ContractIdx = number & { readonly __brand: unique symbol };
+
 /**
  * Represents a Compilation Unit, encapsulating the information necessary for
  * analyzing a single Tact project.
@@ -42,7 +37,7 @@ export class CompilationUnit {
    * Looks for a CFG with a specific index.
    * @returns Found CFG or `undefined` if not found.
    */
-  public findCFGByIdx(idx: BasicBlockIdx): CFG | undefined {
+  public findCFGByIdx(idx: CFGIdx): CFG | undefined {
     const funCfg = this.functions.get(idx);
     if (funCfg) return funCfg;
     return Array.from(this.contracts.values())
@@ -173,6 +168,6 @@ export class Contract {
     public ref: SrcInfo,
     idx: ContractIdx | undefined = undefined,
   ) {
-    this.idx = idx ? idx : IdxGenerator.next("ir_contract");
+    this.idx = idx ? idx : (IdxGenerator.next("ir_contract") as ContractIdx);
   }
 }

src/internals/ir/types.ts

@@ -1,8 +1,3 @@
 export type ProjectName = string & { readonly __brand: unique symbol };
-export type FunctionName = string;
-export type ContractName = string;
-
-export type EdgeIdx = number;
-export type BasicBlockIdx = number;
-export type CFGIdx = number;
-export type ContractIdx = number;
+export type FunctionName = string & { readonly __brand: unique symbol };
+export type ContractName = string & { readonly __brand: unique symbol };