From 4c2c395e323860748e39bdb6d00e16e5bb329b50 Mon Sep 17 00:00:00 2001
From: malatrax <simon.malatrait@grenoble-inp.org>
Date: Fri, 12 Jul 2024 16:39:17 +0200
Subject: [PATCH] refactor: update CairoRunner for compilation artifacts

---
 src/cli.ts                 |  5 ++++-
 src/errors/cairoRunner.ts  |  6 ++++++
 src/runners/cairoRunner.ts | 14 ++++++++++----
 src/scripts/run.ts         |  2 +-
 src/vm/program.test.ts     |  5 +++--
 src/vm/program.ts          | 22 ++++++++++++++++++++--
 6 files changed, 44 insertions(+), 10 deletions(-)

diff --git a/src/cli.ts b/src/cli.ts
index ecdce835..c987ccc9 100755
--- a/src/cli.ts
+++ b/src/cli.ts
@@ -29,8 +29,11 @@ program
       return path;
     })
   )
+  .option(
+    '--cairo <FUNCTION>',
+    'Run a Cairo program, executing the given function'
+  )
   .option('-s, --silent', 'silent all logs')
-  .option('--cairo', 'Run a Cairo program')
   .option('--no-relocate', 'do not relocate memory')
   .addOption(
     new Option(
diff --git a/src/errors/cairoRunner.ts b/src/errors/cairoRunner.ts
index f12adf3f..d5ac72fa 100644
--- a/src/errors/cairoRunner.ts
+++ b/src/errors/cairoRunner.ts
@@ -3,3 +3,9 @@ class CairoRunnerError extends Error {}
 export class EmptyRelocatedMemory extends CairoRunnerError {}
 
 export class CairoZeroHintsNotSupported extends CairoRunnerError {}
+
+export class UndefinedEntrypoint extends CairoRunnerError {
+  constructor(name: string) {
+    super(`The function to be executed doesn't exist: ${name}`);
+  }
+}
diff --git a/src/runners/cairoRunner.ts b/src/runners/cairoRunner.ts
index 3d61125f..cbd25bb7 100644
--- a/src/runners/cairoRunner.ts
+++ b/src/runners/cairoRunner.ts
@@ -3,6 +3,7 @@ import * as fs from 'fs';
 import {
   CairoZeroHintsNotSupported,
   EmptyRelocatedMemory,
+  UndefinedEntrypoint,
 } from 'errors/cairoRunner';
 
 import { Felt } from 'primitives/felt';
@@ -74,12 +75,17 @@ export class CairoRunner {
     return new CairoRunner(program, program.data, mainOffset, builtins);
   }
 
-  static fromCairoProgram(program: CairoProgram): CairoRunner {
+  static fromCairoProgram(
+    program: CairoProgram,
+    fn_name: string = 'main'
+  ): CairoRunner {
+    const fn = program.entry_points_by_function[fn_name];
+    if (!fn) throw new UndefinedEntrypoint(fn_name);
     return new CairoRunner(
       program,
       program.bytecode,
-      program.entrypoint,
-      program.builtins,
+      fn.offset,
+      fn.builtins,
       program.hints
     );
   }
@@ -143,7 +149,7 @@ export class CairoRunner {
   }
 
   getOutput() {
-    const builtins = this.program.builtins;
+    const builtins = (this.program as CairoZeroProgram).builtins;
     const outputIdx = builtins.findIndex((name) => name === 'output');
     return outputIdx >= 0 ? this.vm.memory.segments[outputIdx + 2] : [];
   }
diff --git a/src/scripts/run.ts b/src/scripts/run.ts
index 0c306f3f..c8103218 100644
--- a/src/scripts/run.ts
+++ b/src/scripts/run.ts
@@ -46,7 +46,7 @@ export const run = (
 
     if (cairo) {
       const program = parseCairoProgram(file);
-      runner = CairoRunner.fromCairoProgram(program);
+      runner = CairoRunner.fromCairoProgram(program, cairo);
     } else {
       const program = parseCairoZeroProgram(file);
       runner = CairoRunner.fromCairoZeroProgram(program);
diff --git a/src/vm/program.test.ts b/src/vm/program.test.ts
index 22a9730a..384645fa 100644
--- a/src/vm/program.test.ts
+++ b/src/vm/program.test.ts
@@ -63,8 +63,9 @@ describe('program', () => {
       const program = parseCairoProgram(programContent);
       expect(program.bytecode).toEqual(bytecode);
       expect(program.hints).toEqual(hints);
-      expect(program.entrypoint).toEqual(programJson.entrypoint);
-      expect(program.builtins).toEqual(programJson.builtins);
+      expect(program.entry_points_by_function).toEqual(
+        programJson.entry_points_by_function
+      );
     });
   });
 });
diff --git a/src/vm/program.ts b/src/vm/program.ts
index 677b5424..65bac5b4 100644
--- a/src/vm/program.ts
+++ b/src/vm/program.ts
@@ -40,7 +40,6 @@ export type Identifier = z.infer<typeof identifier>;
 const programBase = z.object({
   prime: z.string(),
   compiler_version: z.string(),
-  builtins: z.array(z.string()),
 });
 
 const cairoZeroProgram = programBase.extend({
@@ -58,12 +57,31 @@ const cairoZeroProgram = programBase.extend({
   reference_manager: referenceManager,
 });
 
+const cairoArg = z.object({
+  generic_id: z.string(),
+  size: z.number(),
+  debug_name: z.string(),
+});
+
+const cairoOutputArg = cairoArg.extend({
+  panic_inner_type: cairoArg.optional(),
+});
+
+const entrypoint = z.object({
+  offset: z.number(),
+  builtins: z.array(z.string()),
+  input_args: z.array(cairoArg),
+  return_arg: z.array(cairoOutputArg),
+});
+
+export type Entrypoint = z.infer<typeof entrypoint>;
+
 const cairoProgram = programBase.extend({
   bytecode: z
     .array(z.string())
     .transform((value) => value.map((v) => new Felt(BigInt(v)))),
   hints,
-  entrypoint: z.number(),
+  entry_points_by_function: z.record(z.string(), entrypoint),
 });
 
 export type CairoZeroProgram = z.infer<typeof cairoZeroProgram>;