Fix warpmemory function naming and implicits + other minor stuff

src/cairoWriter/writers/cairoContractWriter.ts

@@ -24,17 +24,17 @@ export class CairoContractWriter extends CairoASTNodeWriter {
         `// This contract may be abstract, it may not implement an abstract parent's methods\n// completely or it may not invoke an inherited contract's constructor correctly.\n`,
       ];
 
-    const dynamicVariables = [...node.dynamicStorageAllocations.entries()].map(
+    const staticVariables = [...node.staticStorageAllocations.entries()].map(
       ([decl, loc]) => `const ${decl.name}: felt252 = ${loc};`,
     );
-    const staticVariables = [...node.staticStorageAllocations.entries()].map(
+    const dynamicVariables = [...node.dynamicStorageAllocations.entries()].map(
       ([decl, loc]) => `const ${decl.name}: felt252 = ${loc};`,
     );
     const variables = [
-      `// Dynamic variables - Arrays and Maps`,
-      ...dynamicVariables,
-      `// Static variables`,
+      staticVariables.length > 0 ? '// Static variables' : '',
       ...staticVariables,
+      dynamicVariables.length > 0 ? '// Dynamic variables - Arrays and Maps' : '',
+      ...dynamicVariables,
     ];
 
     let documentation = getDocumentation(node.documentation, writer);

src/cairoWriter/writers/sourceUnitWriter.ts

@@ -16,12 +16,12 @@ export class SourceUnitWriter extends CairoASTNodeWriter {
     this.generateInterfaceNameMappings(node);
 
     // Every sourceUnit should only define a single contract
-    const mainContract_ =
+    const mainContracts =
       node.vContracts.length >= 2
         ? node.vContracts.filter((cd) => !cd.name.endsWith(TEMP_INTERFACE_SUFFIX))
         : node.vContracts;
 
-    assert(mainContract_.length <= 1, 'Every SourceUnit should only define a single contract');
+    assert(mainContracts.length <= 1, 'Every SourceUnit should only define a single contract');
 
     structRemappings = new Map();
 

src/passes/cairoUtilImporter.ts

@@ -12,13 +12,16 @@ import {
 } from 'solc-typed-ast';
 import { AST } from '../ast/ast';
 import { ASTMapper } from '../ast/mapper';
+import { CairoFunctionDefinition } from '../export';
 import { requireNonNullish } from '../export';
 import { createImport } from '../utils/importFuncGenerator';
 import {
   INTO,
   U256_FROM_FELTS,
   U128_FROM_FELT,
   CONTRACT_ADDRESS,
+  WARP_MEMORY,
+  WM_INIT,
   CUTOFF_DOWNCAST,
   WARPLIB_INTEGER,
 } from '../utils/importPaths';
@@ -31,16 +34,27 @@ import { getContainingSourceUnit, primitiveTypeToCairo } from '../utils/utils';
   the warplib maths functions as they are added to the code, but for determining if
   Uint256 needs to be imported, it's easier to do it here
 */
-
 export class CairoUtilImporter extends ASTMapper {
   private dummySourceUnit: SourceUnit | undefined;
 
   // Function to add passes that should have been run before this pass
   addInitialPassPrerequisites(): void {
-    const passKeys: Set<string> = new Set<string>([]);
+    // It requires the AnnotateImplicits (An) because we need to know if any functions
+    // uses WarpMemory
+    // It requires Function Prunner (Fp) for it not to delete any of these newly added
+    // imports
+    const passKeys: Set<string> = new Set<string>(['An', 'Fp']);
     passKeys.forEach((key) => this.addPassPrerequisite(key));
   }
 
+  visitCairoFunctionDefinition(node: CairoFunctionDefinition, ast: AST): void {
+    if (node.implicits.has('warp_memory')) {
+      createImport(...WARP_MEMORY, this.dummySourceUnit ?? node, ast);
+      createImport(...WM_INIT, this.dummySourceUnit ?? node, ast);
+    }
+    this.commonVisit(node, ast);
+  }
+
   visitElementaryTypeName(node: ElementaryTypeName, ast: AST): void {
     const cairoType = primitiveTypeToCairo(node.name);
     if (cairoType === 'u256') {
@@ -63,7 +77,7 @@ export class CairoUtilImporter extends ASTMapper {
 
   visitLiteral(node: Literal, ast: AST): void {
     const type = safeGetNodeType(node, ast.inference);
-    if (type instanceof IntType && type.nBits > 251) {
+    if (type instanceof IntType && type.nBits === 256) {
       createImport(...U256_FROM_FELTS, this.dummySourceUnit ?? node, ast);
     }
 

src/passes/references/memoryAllocations.ts

@@ -17,9 +17,13 @@ import { printNode } from '../../utils/astPrinter';
 import { CairoType, TypeConversionContext } from '../../utils/cairoTypeSystem';
 import { NotSupportedYetError } from '../../utils/errors';
 import { createCallToFunction } from '../../utils/functionGeneration';
-import { createNumberLiteral, createUint256TypeName } from '../../utils/nodeTemplates';
+import {
+  createFeltTypeName,
+  createNumberLiteral,
+  createUint256TypeName,
+} from '../../utils/nodeTemplates';
 import { getElementType, safeGetNodeType } from '../../utils/nodeTypeProcessing';
-import { WM_NEW } from '../../utils/importPaths';
+import { U256_TO_FELT252, WM_NEW } from '../../utils/importPaths';
 
 /*
   Handles expressions that directly insert data into memory: struct constructors, news, and inline arrays
@@ -83,8 +87,8 @@ export class MemoryAllocations extends ReferenceSubPass {
       node,
       ...WM_NEW,
       [
-        ['len', createUint256TypeName(ast)],
-        ['elemWidth', createUint256TypeName(ast)],
+        ['len', createFeltTypeName(ast)],
+        ['elem_size', createFeltTypeName(ast)],
       ],
       [['loc', node.vExpression.vTypeName, DataLocation.Memory]],
     );
@@ -102,9 +106,21 @@ export class MemoryAllocations extends ReferenceSubPass {
       TypeConversionContext.Ref,
     );
 
+    const uint256ToFeltDef = ast.registerImport(
+      node,
+      ...U256_TO_FELT252,
+      [['x_256', createUint256TypeName(ast)]],
+      [['x_252', createFeltTypeName(ast)]],
+    );
+    const uint256ToFeltConversion = createCallToFunction(
+      uint256ToFeltDef,
+      [node.vArguments[0]],
+      ast,
+    );
+
     const call = createCallToFunction(
       funcImport,
-      [node.vArguments[0], createNumberLiteral(elementCairoType.width, ast, 'uint256')],
+      [uint256ToFeltConversion, createNumberLiteral(elementCairoType.width, ast, 'uint252')],
       ast,
     );
 

src/utils/cairoTypeSystem.ts

@@ -96,7 +96,7 @@ export abstract class CairoType {
     } else if (tp instanceof FunctionType) {
       throw new NotSupportedYetError('Serialising FunctionType not supported yet');
     } else if (tp instanceof IntType) {
-      return new CairoUint(tp.nBits);
+      return tp.nBits === 252 ? new CairoFelt() : new CairoUint(tp.nBits);
     } else if (tp instanceof MappingType) {
       return new WarpLocation();
     } else if (tp instanceof PointerType) {

src/utils/importFuncGenerator.ts

@@ -91,23 +91,6 @@ const FUNCTION_IMPORTS = [
   Paths.BOOL_INTO_FELT252,
   Paths.FELT252_INTO_BOOL,
   Paths.WARPLIB_INTEGER,
-  Paths.WARP_MEMORY,
-  Paths.WM_NEW,
-  Paths.WM_READ,
-  Paths.WM_ALLOC,
-  Paths.WM_STORE,
-  Paths.WM_WRITE,
-  Paths.WM_CREATE,
-  Paths.WM_GET_ID,
-  Paths.WM_RETRIEVE,
-  Paths.WM_INDEX_DYN,
-  Paths.WM_UNSAFE_READ,
-  Paths.WM_INDEX_STATIC,
-  Paths.WM_UNSAFE_ALLOC,
-  Paths.WM_UNSAFE_WRITE,
-  Paths.WM_READ_MULTIPLE,
-  Paths.WM_WRITE_MULTIPLE,
-  Paths.WM_DYN_ARRAY_LENGTH,
 ];
 
 export function createImport(
@@ -175,7 +158,7 @@ export function createImport(
     return createFuncImport();
   }
 
-  throw new TranspileFailedError(`Import ${name} from ${path} is not defined.`);
+  throw new TranspileFailedError(`Import ${name} from ${path.join('::')} is not defined.`);
 }
 
 function findExistingImport(

src/utils/importPaths.ts

@@ -245,6 +245,8 @@ const MEMORY_MODULE = ['warplib', 'warp_memory'];
 
 export const WARP_MEMORY: [string[], string] = [MEMORY_MODULE, 'WarpMemory'];
 const WARP_MEMORY_TRAIT: string[] = [...MEMORY_MODULE, 'WarpMemoryTrait'];
+
+export const WM_INIT: [string[], string] = [WARP_MEMORY_TRAIT, 'warp_memory.initialize'];
 export const WM_READ: [string[], string] = [WARP_MEMORY_TRAIT, 'warp_memory.read'];
 export const WM_WRITE: [string[], string] = [WARP_MEMORY_TRAIT, 'warp_memory.write'];
 export const WM_ALLOC: [string[], string] = [WARP_MEMORY_TRAIT, 'warp_memory.alloc'];
@@ -261,7 +263,7 @@ export const WM_UNSAFE_ALLOC: [string[], string] = [
   'warp_memory.unsafe_alloc',
 ];
 
-const WARP_MEMORY_ARRAYS_TRAIT: string[] = [...MEMORY_MODULE, 'WarpMemoryArraysTrait'];
+const WARP_MEMORY_ARRAYS_TRAIT: string[] = [...MEMORY_MODULE, 'arrays', 'WarpMemoryArraysTrait'];
 export const WM_DYN_ARRAY_LENGTH: [string[], string] = [
   WARP_MEMORY_ARRAYS_TRAIT,
   'warp_memory.length_dyn',
@@ -280,7 +282,11 @@ export const WM_GET_ID: [string[], string] = [
   'warp_memory.get_or_create_id',
 ];
 
-const MULTICELL_ACCESSOR_TRAIT: string[] = [...MEMORY_MODULE, 'WarpMemoryMultiCellAccessorTrait'];
+const MULTICELL_ACCESSOR_TRAIT: string[] = [
+  ...MEMORY_MODULE,
+  'accessors',
+  'WarpMemoryMultiCellAccessorTrait',
+];
 export const WM_WRITE_MULTIPLE: [string[], string] = [
   MULTICELL_ACCESSOR_TRAIT,
   'warp_memory.write_multiple',

src/warplib/gatherWarplibImports.ts

@@ -3,6 +3,7 @@ import { Implicits } from '../utils/utils';
 import { parseMultipleRawCairoFunctions } from '../utils/cairoParsing';
 import { glob } from 'glob';
 import path from 'path';
+import assert from 'assert';
 
 import { WARP_ROOT } from '../config';
 
@@ -18,6 +19,23 @@ export const warplibImportInfo = glob
         .slice('warplib/src/'.length, -'.cairo'.length)
         .split(path.sep),
     ];
+
+    // Handle WarpMemory function gathering differently
+    if (importPath.includes('warp_memory')) {
+      // Hardcode WarpMemory struct
+      if (pathToFile.endsWith('warp_memory.cairo')) {
+        warplibMap.set(importPath.join('/'), new Map([['WarpMemory', []]]));
+      }
+      gatherWarpMemoryFuncs(rawCairoCode, importPath, warplibMap);
+      return warplibMap;
+    }
+
+    if (importPath[importPath.length - 1] === 'integer') {
+      console.log('Gather int conversion');
+      gatherIntegerConversion(rawCairoCode, importPath, warplibMap);
+      return warplibMap;
+    }
+
     // TODO: Add encodePath here. Importing encodePath cause circular
     // dependency error. Suggested solution is to relocate all import
     // related scripts (including this, and the ones in src/utils)
@@ -32,3 +50,91 @@ export const warplibImportInfo = glob
     );
     return warplibMap;
   }, new Map<string, Map<string, Implicits[]>>());
+
+function gatherIntegerConversion(
+  rawCairoCode: string,
+  importPath: string[],
+  warplibMap: Map<string, Map<string, Implicits[]>>,
+) {
+  const integerTraitInfo = traitFunctionExtractor(rawCairoCode);
+  console.log(integerTraitInfo);
+
+  integerTraitInfo.forEach(({ trait, code }) => {
+    const traitImportPath = [...importPath, trait];
+    const key = traitImportPath.join('/');
+
+    const fileMap: Map<string, Implicits[]> = warplibMap.get(key) ?? new Map<string, Implicits[]>();
+    if (!warplibMap.has(key)) {
+      warplibMap.set(key, fileMap);
+    }
+
+    parseMultipleRawCairoFunctions(code).forEach((cairoFunc) => fileMap.set(cairoFunc.name, []));
+  });
+}
+
+// Functions that stored warp memory functions accordingly
+function gatherWarpMemoryFuncs(
+  rawCairoCode: string,
+  importPath: string[],
+  warplibMap: Map<string, Map<string, Implicits[]>>,
+) {
+  // Get the body for each trait
+  const warpMemoryTraitInfo = traitFunctionExtractor(rawCairoCode);
+
+  // Store each trait import path with each function prefixed with "warp_memory."
+  warpMemoryTraitInfo.forEach(({ trait, code }) => {
+    const traitImportPath = [...importPath, trait];
+    const key = traitImportPath.join('/');
+
+    const fileMap: Map<string, Implicits[]> = warplibMap.get(key) ?? new Map<string, Implicits[]>();
+    if (!warplibMap.has(key)) {
+      warplibMap.set(key, fileMap);
+    }
+
+    parseMultipleRawCairoFunctions(code).forEach((cairoFunc) =>
+      fileMap.set(`warp_memory.${cairoFunc.name}`, ['warp_memory']),
+    );
+  });
+}
+
+type TraitInfo = { trait: string; code: string };
+function traitFunctionExtractor(rawCairoCode: string): TraitInfo[] {
+  const warpMemoryInfo: TraitInfo[] = [];
+  let currentTrait = '';
+  let currentBody = '';
+  rawCairoCode
+    .split('\n')
+    .filter((l) => !l.match('^[ ]*$'))
+    .forEach((l) => {
+      if (l.trim().startsWith('trait')) {
+        if (currentTrait !== '') {
+          warpMemoryInfo.push({ trait: currentTrait, code: currentBody });
+        }
+        currentTrait = extractName('trait', l);
+        currentBody = '';
+        return;
+      }
+      if (l.trim().startsWith('impl')) {
+        if (currentTrait !== '') {
+          warpMemoryInfo.push({ trait: currentTrait, code: currentBody });
+        }
+        currentTrait = '';
+        currentBody = '';
+        return;
+      }
+      currentBody += '\n' + l.trim();
+    });
+
+  return warpMemoryInfo;
+}
+
+function extractName(keyword: string, line: string) {
+  const regex = new RegExp(`${keyword}[ ]+(?<name>\\w+)`);
+  const m = line.match(regex);
+  assert(
+    m !== null && m.groups !== undefined,
+    `Error extracting name for ${keyword} in line:\n"${line}"`,
+  );
+
+  return m.groups.name;
+}