Fix environment restoration and implement compound functions - Add RESTORE_ENV instruction to properly restore environment after function calls - Fix function parameter binding in CSE machine - Compound functions now work correctly (functions using other functions) - Linear recursion and iteration working properly

Author: TRK95

src/CSE-machine/astToControl.ts

@@ -6,7 +6,7 @@ import {
 } from "../transpiler/types/nodes/scheme-node-types";
 
 export function astToControl(expr: Expression): ControlItem[] {
-  // Hỗ trợ cả node Atomic và Extended
+ 
   if (
     expr instanceof Atomic.NumericLiteral ||
     expr instanceof Atomic.BooleanLiteral ||

src/CSE-machine/instrCreator.ts

@@ -5,6 +5,7 @@ import {
   Extended,
 } from "../transpiler/types/nodes/scheme-node-types";
 import { Location, Position } from "../transpiler/types/location";
+import { Environment } from "./environment";
 import {
   Instr,
   InstrType,
@@ -22,6 +23,7 @@ import {
   CarInstr,
   CdrInstr,
   ConsInstr,
+  RestoreEnvInstr,
   AppInstr,
   BranchInstr,
 } from "./types";
@@ -187,6 +189,14 @@ export function createBranchInstr(
     alternate,
   };
 }
+
+export function createRestoreEnvInstr(env: Environment): RestoreEnvInstr {
+  return {
+    instrType: InstrType.RESTORE_ENV,
+    srcNode: { type: "StatementSequence", body: [], location: { start: { line: 0, column: 0 }, end: { line: 0, column: 0 } } },
+    env,
+  };
+}
 import { LiteralInstr, VariableInstr, LambdaInstr, IfInstr } from "./types";
 
 // Literal instruction

src/CSE-machine/interpreter.ts

@@ -31,6 +31,7 @@ import {
   CarInstr,
   CdrInstr,
   ConsInstr,
+  RestoreEnvInstr,
   AppInstr,
   BranchInstr,
   StatementSequence,
@@ -197,12 +198,10 @@ function evaluateExpression(
     }
   } else if (expr instanceof Atomic.Conditional) {
     console.log("DEBUG: Evaluating Conditional expression");
-    // Push branch instruction AFTER test evaluation
-    // This ensures test is evaluated and pushed to stash before branch runs
+    // Push test expression first, then branch instruction
+    // The branch instruction will decide which branch to take based on test result
     control.push(instr.createBranchInstr(expr.consequent, expr.alternate));
     control.push(expr.test);
-    control.push(expr.consequent);
-    control.push(expr.alternate);
   } else if (expr instanceof Atomic.Lambda) {
     // Create closure
     const closure: Value = {
@@ -314,12 +313,22 @@ function evaluateInstruction(
       console.log("DEBUG: Arguments:", args);
 
       if (operator.type === "closure") {
-        // Apply closure
+        // Apply closure - save current environment and create new one
+        const currentEnv = context.environment;
         const newEnv = createBlockEnvironment(operator.env);
+        
+        // Bind parameters to the new environment
         for (let i = 0; i < operator.params.length; i++) {
           newEnv.define(operator.params[i], args[i] || { type: "nil" });
         }
+        
+        // Set the new environment for function execution
         context.environment = newEnv;
+        
+        // Push a marker to restore environment after function execution
+        control.push(instr.createRestoreEnvInstr(currentEnv));
+        
+        // Push function body for execution
         control.push(...operator.body);
       } else if (operator.type === "primitive") {
         // Apply primitive function
@@ -349,7 +358,9 @@ function evaluateInstruction(
       console.log("DEBUG: Test value:", test);
       const branchInstr = instruction as BranchInstr;
 
-      if (test.type === "boolean" && test.value) {
+      // Check if test is truthy (not false and not nil)
+      const isTruthy = test.type !== "boolean" || test.value !== false;
+      if (isTruthy && test.type !== "nil") {
         console.log("DEBUG: Taking consequent branch");
         control.push(branchInstr.consequent);
       } else if (branchInstr.alternate) {
@@ -454,6 +465,12 @@ function evaluateInstruction(
       break;
     }
 
+    case InstrType.RESTORE_ENV: {
+      const restoreInstr = instruction as RestoreEnvInstr;
+      context.environment = restoreInstr.env;
+      break;
+    }
+
     default:
       throw new Error(`Unsupported instruction type: ${instruction.instrType}`);
   }

src/CSE-machine/types.ts

@@ -60,6 +60,7 @@ export enum InstrType {
   CAR = "Car",
   CDR = "Cdr",
   CONS = "Cons",
+  RESTORE_ENV = "RestoreEnv",
 }
 
 interface BaseInstr {
@@ -175,6 +176,10 @@ export interface ConsInstr extends BaseInstr {
   cdr: Expression;
 }
 
+export interface RestoreEnvInstr extends BaseInstr {
+  env: Environment;
+}
+
 // ...existing code...
 
 export interface LiteralInstr extends BaseInstr {
@@ -219,6 +224,7 @@ export type Instr =
   | CarInstr
   | CdrInstr
   | ConsInstr
+  | RestoreEnvInstr
   | UnOpInstr
   | BinOpInstr
   | LiteralInstr

src/test/12-test-chapter1-features.ts

@@ -220,6 +220,112 @@ function testChapter1Features() {
       expected: { type: "number", value: 3 },
       description: "Length of variable list",
     },
+
+    // ===== COMPOUND FUNCTIONS =====
+    {
+      code: "(define (square x) (* x x))",
+      expected: { type: "void" },
+      description: "Define square function",
+    },
+    {
+      code: "(define (cube x) (* x (square x)))",
+      expected: { type: "void" },
+      description: "Define cube using square (compound function)",
+    },
+    {
+      code: "(cube 3)",
+      expected: { type: "number", value: 27 },
+      description: "Test compound function",
+    },
+    {
+      code: "(define (average x y) (/ (+ x y) 2))",
+      expected: { type: "void" },
+      description: "Define average function",
+    },
+    {
+      code: "(define (square-average x y) (square (average x y)))",
+      expected: { type: "void" },
+      description: "Compound function using average",
+    },
+    {
+      code: "(square-average 4 6)",
+      expected: { type: "number", value: 25 },
+      description: "Test nested compound function",
+    },
+
+    // ===== LINEAR RECURSION =====
+    {
+      code: "(define (factorial n) (if (= n 0) 1 (* n (factorial (- n 1)))))",
+      expected: { type: "void" },
+      description: "Define factorial with linear recursion",
+    },
+    {
+      code: "(factorial 5)",
+      expected: { type: "number", value: 120 },
+      description: "Test factorial recursion",
+    },
+    {
+      code: "(factorial 0)",
+      expected: { type: "number", value: 1 },
+      description: "Test factorial base case",
+    },
+    {
+      code: "(define (sum-list lst) (if (null? lst) 0 (+ (car lst) (sum-list (cdr lst)))))",
+      expected: { type: "void" },
+      description: "Define sum-list with linear recursion",
+    },
+    {
+      code: "(sum-list (list 1 2 3 4 5))",
+      expected: { type: "number", value: 15 },
+      description: "Test sum-list recursion",
+    },
+    {
+      code: "(sum-list ())",
+      expected: { type: "number", value: 0 },
+      description: "Test sum-list base case",
+    },
+    {
+      code: "(define (count-elements lst) (if (null? lst) 0 (+ 1 (count-elements (cdr lst)))))",
+      expected: { type: "void" },
+      description: "Define count-elements with linear recursion",
+    },
+    {
+      code: "(count-elements (list 'a 'b 'c))",
+      expected: { type: "number", value: 3 },
+      description: "Test count-elements recursion",
+    },
+
+    // ===== ITERATION (Tail Recursion) =====
+    {
+      code: "(define (factorial-iter n) (define (iter n acc) (if (= n 0) acc (iter (- n 1) (* n acc)))) (iter n 1))",
+      expected: { type: "void" },
+      description: "Define factorial with iteration (tail recursion)",
+    },
+    {
+      code: "(factorial-iter 5)",
+      expected: { type: "number", value: 120 },
+      description: "Test factorial iteration",
+    },
+    {
+      code: "(define (sum-list-iter lst) (define (iter lst acc) (if (null? lst) acc (iter (cdr lst) (+ (car lst) acc)))) (iter lst 0))",
+      expected: { type: "void" },
+      description: "Define sum-list with iteration",
+    },
+    {
+      code: "(sum-list-iter (list 1 2 3 4 5))",
+      expected: { type: "number", value: 15 },
+      description: "Test sum-list iteration",
+    },
+    {
+      code: "(define (reverse-list lst) (define (iter lst result) (if (null? lst) result (iter (cdr lst) (cons (car lst) result)))) (iter lst ()))",
+      expected: { type: "void" },
+      description: "Define reverse-list with iteration",
+    },
+    {
+      code: "(reverse-list (list 1 2 3))",
+      expected: { type: "list", elements: [{ type: "number", value: 3 }, { type: "number", value: 2 }, { type: "number", value: 1 }] },
+      description: "Test reverse-list iteration",
+    },
   ];
 
   let passed = 0;
@@ -287,6 +393,9 @@ function testChapter1Features() {
   console.log("✅ Built-in functions with variables");
   console.log("✅ Boolean operations with variables");
   console.log("✅ Function composition");
+  console.log("✅ Compound functions (functions using other functions)");
+  console.log("✅ Linear recursion (factorial, sum-list, count-elements)");
+  console.log("✅ Iteration via tail recursion (factorial-iter, sum-list-iter, reverse-list)");
 }
 
 testChapter1Features();