zexus 1.6.8 → 1.7.2

package/src/zexus/vm/optimizer.py

@@ -70,6 +70,75 @@ class BytecodeOptimizer:
         self.max_passes = max_passes
         self.debug = debug
         self.stats = OptimizationStats()
+        self.min_basic_size = 2
+        self.min_aggressive_size = 2
+        self.min_experimental_size = 20
+        self._control_flow_ops = {
+            "JUMP",
+            "JUMP_IF_FALSE",
+            "JUMP_IF_TRUE",
+            "JUMP_FORWARD",
+            "JUMP_BACKWARD",
+            "LABEL",
+            "JUMP_TARGET",
+        }
+
+    def _has_control_flow(self, instructions: List[Tuple[str, Any]]) -> bool:
+        """Check if instruction stream contains control-flow operations."""
+        for op, _ in instructions:
+            if op in self._control_flow_ops:
+                return True
+        return False
+
+    def _validate_control_flow(self, instructions: List[Tuple[str, Any]]) -> bool:
+        """Validate that jump targets are sane and within bounds."""
+        max_index = len(instructions) - 1
+        for idx, (op, operand) in enumerate(instructions):
+            if op in ("JUMP", "JUMP_IF_FALSE", "JUMP_IF_TRUE", "JUMP_FORWARD", "JUMP_BACKWARD"):
+                if not isinstance(operand, int):
+                    return False
+                if operand < 0 or operand > max_index:
+                    return False
+                if op == "JUMP" and operand == idx:
+                    return False
+                if op == "JUMP_FORWARD" and operand <= idx:
+                    return False
+                if op == "JUMP_BACKWARD" and operand >= idx:
+                    return False
+        return True
+
+    def _run_pass(
+        self,
+        name: str,
+        instructions: List[Tuple[str, Any]],
+        apply_fn,
+        *,
+        can_change_size: bool = True
+    ) -> List[Tuple[str, Any]]:
+        """Run a pass with logging and control-flow validation."""
+        if can_change_size and self._has_control_flow(instructions):
+            if self.debug:
+                print(f"[Optimizer] {name}: skipped (control flow present)")
+            return instructions
+
+        before_len = len(instructions)
+        try:
+            optimized = apply_fn()
+        except Exception as exc:
+            if self.debug:
+                print(f"[Optimizer] {name}: failed ({exc})")
+            return instructions
+
+        after_len = len(optimized)
+        if self.debug:
+            print(f"[Optimizer] {name}: {before_len} -> {after_len}")
+
+        if not self._validate_control_flow(optimized):
+            if self.debug:
+                print(f"[Optimizer] {name}: invalid control flow detected, reverting")
+            return instructions
+
+        return optimized
     
     def optimize(self, instructions: List[Tuple[str, Any]], constants: List[Any] = None) -> List[Tuple[str, Any]]:
         """
@@ -97,21 +166,75 @@ class BytecodeOptimizer:
             
             # Level 1: Basic optimizations
             if self.level >= 1:
-                optimized = self._constant_folding(optimized, constants)
-                optimized = self._dead_code_elimination(optimized)
-                optimized = self._peephole_optimization(optimized)
+                if len(optimized) >= self.min_basic_size:
+                    optimized = self._run_pass(
+                        "constant_folding",
+                        optimized,
+                        lambda: self._constant_folding(optimized, constants),
+                        can_change_size=True,
+                    )
+                    optimized = self._run_pass(
+                        "dead_code_elimination",
+                        optimized,
+                        lambda: self._dead_code_elimination(optimized),
+                        can_change_size=False,
+                    )
+                    optimized = self._run_pass(
+                        "peephole_optimization",
+                        optimized,
+                        lambda: self._peephole_optimization(optimized),
+                        can_change_size=True,
+                    )
+                elif self.debug:
+                    print("[Optimizer] basic passes skipped (small instruction count)")
             
             # Level 2: Aggressive optimizations
             if self.level >= 2:
-                optimized = self._copy_propagation(optimized)
-                optimized = self._instruction_combining(optimized, constants)
-                optimized = self._jump_threading(optimized)
-                optimized = self._strength_reduction(optimized)
+                if len(optimized) >= self.min_aggressive_size:
+                    optimized = self._run_pass(
+                        "copy_propagation",
+                        optimized,
+                        lambda: self._copy_propagation(optimized),
+                        can_change_size=False,
+                    )
+                    optimized = self._run_pass(
+                        "instruction_combining",
+                        optimized,
+                        lambda: self._instruction_combining(optimized, constants),
+                        can_change_size=True,
+                    )
+                    optimized = self._run_pass(
+                        "jump_threading",
+                        optimized,
+                        lambda: self._jump_threading(optimized),
+                        can_change_size=False,
+                    )
+                    optimized = self._run_pass(
+                        "strength_reduction",
+                        optimized,
+                        lambda: self._strength_reduction(optimized),
+                        can_change_size=False,
+                    )
+                elif self.debug:
+                    print("[Optimizer] aggressive passes skipped (small instruction count)")
             
             # Level 3: Experimental optimizations
             if self.level >= 3:
-                optimized = self._common_subexpression_elimination(optimized)
-                optimized = self._loop_invariant_code_motion(optimized)
+                if len(optimized) >= self.min_experimental_size:
+                    optimized = self._run_pass(
+                        "common_subexpression_elimination",
+                        optimized,
+                        lambda: self._common_subexpression_elimination(optimized),
+                        can_change_size=True,
+                    )
+                    optimized = self._run_pass(
+                        "loop_invariant_code_motion",
+                        optimized,
+                        lambda: self._loop_invariant_code_motion(optimized),
+                        can_change_size=True,
+                    )
+                elif self.debug:
+                    print("[Optimizer] experimental passes skipped (small instruction count)")
             
             self.stats.passes_applied += 1
             
@@ -282,24 +405,51 @@ class BytecodeOptimizer:
     def _dead_code_elimination(self, instructions: List[Tuple[str, Any]]) -> List[Tuple[str, Any]]:
         """
         Remove unreachable code after RETURN, unconditional JUMP, etc.
+        Jump targets are remapped after dead code removal to keep control flow valid.
         """
         result = []
         in_dead_code = False
+        jump_targets: Set[int] = set()
+        removed_indices: Set[int] = set()
+
+        for idx, (op, operand) in enumerate(instructions):
+            if op in ("JUMP", "JUMP_IF_FALSE", "JUMP_IF_TRUE", "JUMP_FORWARD", "JUMP_BACKWARD"):
+                if isinstance(operand, int):
+                    jump_targets.add(operand)
         
-        for op, operand in instructions:
+        for idx, (op, operand) in enumerate(instructions):
             if in_dead_code:
                 # Skip until we hit a jump target or label
-                if op in ("LABEL", "JUMP_TARGET"):
+                if op in ("LABEL", "JUMP_TARGET") or idx in jump_targets:
                     in_dead_code = False
                     result.append((op, operand))
                 else:
                     self.stats.dead_code_removed += 1
+                    removed_indices.add(idx)
             else:
                 result.append((op, operand))
                 # Mark dead code after unconditional control flow
                 if op in ("RETURN", "JUMP"):
                     in_dead_code = True
         
+        # Remap jump targets after dead code removal
+        if removed_indices:
+            index_map: Dict[int, int] = {}
+            new_idx = 0
+            for old_idx in range(len(instructions)):
+                if old_idx not in removed_indices:
+                    index_map[old_idx] = new_idx
+                    new_idx += 1
+            
+            jump_ops = ("JUMP", "JUMP_IF_FALSE", "JUMP_IF_TRUE", "JUMP_FORWARD", "JUMP_BACKWARD")
+            remapped = []
+            for op, operand in result:
+                if op in jump_ops and isinstance(operand, int) and operand in index_map:
+                    remapped.append((op, index_map[operand]))
+                else:
+                    remapped.append((op, operand))
+            result = remapped
+        
         return result
     
     def _peephole_optimization(self, instructions: List[Tuple[str, Any]]) -> List[Tuple[str, Any]]:

package/src/zexus/vm/parallel_vm.py

@@ -28,6 +28,8 @@ import time
 import logging
 from collections import defaultdict
 import traceback
+import signal
+import concurrent.futures
 
 from .bytecode import Bytecode, Opcode
 from .vm import VM
@@ -111,13 +113,17 @@ def _execute_chunk_helper(args):
     Returns:
         ExecutionResult with execution status and metrics
     """
-    chunk, shared_state_dict, retry_count = args
+    chunk, shared_state_dict, retry_count, builtins, parent_env = args
     
     try:
         start_time = time.time()
         
         # Create a minimal VM for this worker
         vm = VM()
+        if builtins is not None:
+            vm.builtins = builtins
+        if parent_env is not None:
+            vm._parent_env = parent_env
         
         # Load shared state
         for var, value in shared_state_dict.items():
@@ -125,6 +131,7 @@ def _execute_chunk_helper(args):
         
         # Create bytecode from chunk
         bytecode = Bytecode()
+        bytecode.constants = list(chunk.constants)
         for opcode, operand in chunk.instructions:
             bytecode.instructions.append((opcode, operand))
         
@@ -184,6 +191,7 @@ class BytecodeChunk:
     variables_read: Set[str] = field(default_factory=set)
     variables_written: Set[str] = field(default_factory=set)
     can_parallelize: bool = True
+    constants: List[Any] = field(default_factory=list)
     
     def __repr__(self):
         ins_count = len(self.instructions)
@@ -302,6 +310,7 @@ class BytecodeChunker:
                         chunk.variables_read.update(reads)
                         chunk.variables_written.update(writes)
             
+            chunk.constants = list(bytecode.constants)
             chunks.append(chunk)
         
         # Detect dependencies
@@ -398,8 +407,14 @@ class ResultMerger:
         for result in sorted_results:
             merged_variables.update(result.variables_modified)
         
-        # Last result is the final result
-        final_result = sorted_results[-1].result if sorted_results else None
+        # Prefer the last non-None result; fall back to merged variables
+        final_result = None
+        for result in reversed(sorted_results):
+            if result.result is not None:
+                final_result = result.result
+                break
+        if final_result is None and merged_variables:
+            final_result = merged_variables
         
         return True, final_result, merged_variables
     
@@ -436,7 +451,12 @@ class WorkerPool:
     def start(self) -> None:
         """Start the worker pool"""
         if self.pool is None:
-            self.pool = Pool(processes=self.num_workers)
+            try:
+                self.pool = Pool(processes=self.num_workers)
+            except (OSError, RuntimeError) as e:
+                logger.warning(f"Failed to start multiprocessing Pool: {e}")
+                self.pool = None
+                raise
     
     def shutdown(self) -> None:
         """Shutdown the worker pool"""
@@ -445,7 +465,13 @@ class WorkerPool:
             self.pool.join()
             self.pool = None
     
-    def execute_chunk(self, chunk: BytecodeChunk, shared_state: SharedState) -> ExecutionResult:
+    def execute_chunk(
+        self,
+        chunk: BytecodeChunk,
+        shared_state: SharedState,
+        builtins: Optional[Dict[str, Any]] = None,
+        parent_env: Optional[Dict[str, Any]] = None,
+    ) -> ExecutionResult:
         """
         Execute a single chunk (called by worker process).
         
@@ -456,6 +482,10 @@ class WorkerPool:
         try:
             # Create a local VM for this chunk
             vm = VM()
+            if builtins is not None:
+                vm.builtins = builtins
+            if parent_env is not None:
+                vm._parent_env = parent_env
             
             # Load shared variables
             for var in chunk.variables_read:
@@ -500,10 +530,14 @@ class WorkerPool:
                 execution_time=execution_time
             )
     
-    def submit_chunks(self, 
-                     chunks: List[BytecodeChunk], 
-                     shared_state: SharedState,
-                     config: ParallelConfig) -> List[ExecutionResult]:
+    def submit_chunks(
+        self,
+        chunks: List[BytecodeChunk],
+        shared_state: SharedState,
+        config: ParallelConfig,
+        builtins: Optional[Dict[str, Any]] = None,
+        parent_env: Optional[Dict[str, Any]] = None,
+    ) -> List[ExecutionResult]:
         """
         Submit chunks for parallel execution with retry logic.
         
@@ -537,38 +571,55 @@ class WorkerPool:
             
             # Execute parallel chunks with retry logic
             if parallel_chunks and len(parallel_chunks) > 1:
-                # Use cloudpickle for better serialization
+                # Use snapshot of shared state for this level
                 shared_dict = dict(shared_state.variables)
-                
-                chunk_results = []
-                for chunk in parallel_chunks:
-                    retry_count = 0
-                    success = False
-                    
-                    while retry_count < config.retry_attempts and not success:
+
+                pending: Dict[int, Tuple[BytecodeChunk, int]] = {
+                    c.chunk_id: (c, 0) for c in parallel_chunks
+                }
+                chunk_results: List[ExecutionResult] = []
+
+                while pending:
+                    futures: Dict[int, Any] = {}
+                    for chunk_id, (chunk, retry_count) in pending.items():
+                        futures[chunk_id] = self.pool.apply_async(
+                            _execute_chunk_helper,
+                            ((chunk, shared_dict, retry_count, builtins, parent_env),)
+                        )
+
+                    next_pending: Dict[int, Tuple[BytecodeChunk, int]] = {}
+                    for chunk_id, future in futures.items():
+                        chunk, retry_count = pending[chunk_id]
                         try:
-                            # Submit with timeout
-                            future = self.pool.apply_async(
-                                _execute_chunk_helper,
-                                ((chunk, shared_dict, retry_count),)
-                            )
                             result = future.get(timeout=config.timeout_seconds)
-                            
                             if result.success:
-                                success = True
                                 chunk_results.append(result)
                                 metrics.chunks_succeeded += 1
                             else:
                                 retry_count += 1
-                                metrics.chunks_retried += 1
-                                logger.warning(f"Chunk {chunk.chunk_id} failed, retry {retry_count}/{config.retry_attempts}")
-                        
+                                if retry_count < config.retry_attempts:
+                                    metrics.chunks_retried += 1
+                                    logger.warning(
+                                        f"Chunk {chunk.chunk_id} failed, retry {retry_count}/{config.retry_attempts}"
+                                    )
+                                    next_pending[chunk_id] = (chunk, retry_count)
+                                else:
+                                    result.retry_count = retry_count
+                                    chunk_results.append(result)
+                                    metrics.chunks_failed += 1
+                                    metrics.errors.append(
+                                        f"Chunk {chunk.chunk_id}: {result.error or 'unknown error'}"
+                                    )
+
                         except mp.TimeoutError:
                             retry_count += 1
-                            metrics.chunks_retried += 1
-                            logger.error(f"Chunk {chunk.chunk_id} timed out after {config.timeout_seconds}s")
-                            
-                            if retry_count >= config.retry_attempts:
+                            if retry_count < config.retry_attempts:
+                                metrics.chunks_retried += 1
+                                logger.error(
+                                    f"Chunk {chunk.chunk_id} timed out after {config.timeout_seconds}s"
+                                )
+                                next_pending[chunk_id] = (chunk, retry_count)
+                            else:
                                 error_result = ExecutionResult(
                                     chunk_id=chunk.chunk_id,
                                     success=False,
@@ -578,13 +629,14 @@ class WorkerPool:
                                 chunk_results.append(error_result)
                                 metrics.chunks_failed += 1
                                 metrics.errors.append(f"Chunk {chunk.chunk_id} timeout")
-                        
+
                         except Exception as e:
                             retry_count += 1
-                            metrics.chunks_retried += 1
-                            logger.error(f"Chunk {chunk.chunk_id} error: {e}")
-                            
-                            if retry_count >= config.retry_attempts:
+                            if retry_count < config.retry_attempts:
+                                metrics.chunks_retried += 1
+                                logger.error(f"Chunk {chunk.chunk_id} error: {e}")
+                                next_pending[chunk_id] = (chunk, retry_count)
+                            else:
                                 error_result = ExecutionResult(
                                     chunk_id=chunk.chunk_id,
                                     success=False,
@@ -595,9 +647,11 @@ class WorkerPool:
                                 chunk_results.append(error_result)
                                 metrics.chunks_failed += 1
                                 metrics.errors.append(f"Chunk {chunk.chunk_id}: {str(e)}")
-                
+
+                    pending = next_pending
+
                 results.extend(chunk_results)
-                
+
                 # Update shared state with results
                 for result in chunk_results:
                     if result.success:
@@ -605,7 +659,7 @@ class WorkerPool:
             
             # Execute sequential chunks one by one
             for chunk in sequential_chunks + (parallel_chunks if len(parallel_chunks) == 1 else []):
-                result = self.execute_chunk(chunk, shared_state)
+                result = self.execute_chunk(chunk, shared_state, builtins=builtins, parent_env=parent_env)
                 results.append(result)
                 
                 if result.success:
@@ -701,11 +755,21 @@ class ParallelVM:
         # Metrics
         self.last_metrics: Optional[ExecutionMetrics] = None
         self.cumulative_metrics = ExecutionMetrics()
+        self.stats = {
+            "parallel_executions": 0,
+            "sequential_executions": 0,
+            "fallback_executions": 0,
+        }
         
         logger.info(f"ParallelVM initialized: {self.config.worker_count} workers, "
                    f"chunk_size={self.config.chunk_size}, mode={mode.value}")
     
-    def execute(self, bytecode: Bytecode, sequential_fallback: bool = True) -> Any:
+    def execute(
+        self,
+        bytecode: Bytecode,
+        sequential_fallback: bool = True,
+        initial_state: Optional[Dict[str, Any]] = None,
+    ) -> Any:
         """
         Execute bytecode in parallel with metrics and error handling.
         
@@ -751,15 +815,41 @@ class ParallelVM:
             metrics.chunk_count = len(chunks)
             logger.info(f"Created {len(chunks)} chunks in {time.time() - chunk_start:.4f}s")
             
-            # Initialize shared state
-            manager = Manager()
-            self.shared_state = SharedState(manager)
+            # Initialize shared state — use thread-safe dict instead of 
+            # multiprocessing.Manager which can hang in container environments
+            try:
+                manager = Manager()
+                self.shared_state = SharedState(manager)
+            except (OSError, RuntimeError) as e:
+                logger.warning(f"multiprocessing.Manager() failed ({e}), falling back to sequential")
+                if sequential_fallback and self.config.enable_fallback:
+                    self.stats["fallback_executions"] += 1
+                    result = self._execute_sequential(bytecode)
+                    metrics.total_time = time.time() - start_time
+                    self.last_metrics = metrics
+                    return result
+                raise
             self.merger = ResultMerger()
+
+            builtins = None
+            parent_env = None
+            if initial_state:
+                env_state = initial_state.get("env")
+                if isinstance(env_state, dict):
+                    self.shared_state.batch_write(env_state)
+                builtins = initial_state.get("builtins")
+                parent_env = initial_state.get("parent_env")
             
             # Execute chunks in parallel
             parallel_start = time.time()
             with self.worker_pool as pool:
-                results = pool.submit_chunks(chunks, self.shared_state, self.config)
+                results = pool.submit_chunks(
+                    chunks,
+                    self.shared_state,
+                    self.config,
+                    builtins=builtins,
+                    parent_env=parent_env,
+                )
             metrics.parallel_time = time.time() - parallel_start
             
             logger.info(f"Parallel execution completed in {metrics.parallel_time:.4f}s")
@@ -788,6 +878,7 @@ class ParallelVM:
             
             if not success and sequential_fallback:
                 logger.warning(f"Parallel execution failed: {final_result}. Falling back to sequential.")
+                self.stats["fallback_executions"] += 1
                 result = self._execute_sequential(bytecode)
                 metrics.total_time = time.time() - start_time
                 self.last_metrics = metrics
@@ -811,6 +902,7 @@ class ParallelVM:
             self.cumulative_metrics.chunks_retried += metrics.chunks_retried
             
             self.last_metrics = metrics
+            self.stats["parallel_executions"] += 1
             
             if self.config.enable_metrics:
                 logger.info(f"Execution metrics: {metrics.to_dict()}")
@@ -826,6 +918,7 @@ class ParallelVM:
             
             if sequential_fallback and self.config.enable_fallback:
                 logger.warning("Falling back to sequential execution due to error")
+                self.stats["fallback_executions"] += 1
                 return self._execute_sequential(bytecode)
             else:
                 raise
@@ -834,7 +927,9 @@ class ParallelVM:
         """Execute bytecode sequentially (fallback mode)"""
         logger.info("Executing in sequential mode")
         vm = VM()
-        return vm.execute(bytecode)
+        result = vm.execute(bytecode)
+        self.stats["sequential_executions"] += 1
+        return result
     
     def get_statistics(self) -> Dict[str, Any]:
         """Get execution statistics"""

package/src/zexus/vm/peephole_optimizer.py

@@ -82,6 +82,8 @@ class PeepholeOptimizer:
     def __init__(self, level: OptimizationLevel = OptimizationLevel.BASIC):
         self.level = level
         self.stats = OptimizationStats()
+        self._consts: Optional[List[Any]] = None
+        self._const_indexed = False
         
         # Optimization pattern matchers
         self.patterns = {
@@ -144,6 +146,63 @@ class PeepholeOptimizer:
         )
         
         return optimized
+
+    def optimize_bytecode(
+        self,
+        instructions: List[Tuple[Any, Any]],
+        constants: Optional[List[Any]] = None
+    ) -> Tuple[List[Tuple[Any, Any]], List[Any]]:
+        """
+        Optimize raw bytecode instructions with constant pool awareness.
+
+        Args:
+            instructions: List of (opcode, operand) tuples
+            constants: Bytecode constants list (optional)
+
+        Returns:
+            (optimized_instructions, updated_constants)
+        """
+        control_flow_ops = {
+            "JUMP",
+            "JUMP_IF_FALSE",
+            "JUMP_IF_TRUE",
+            "JUMP_FORWARD",
+            "JUMP_BACKWARD",
+            "JUMP_TARGET",
+            "LABEL",
+        }
+        for instr in instructions:
+            if instr is None:
+                continue
+            if isinstance(instr, tuple) and len(instr) >= 1:
+                op = instr[0]
+                op_name = op.name if hasattr(op, "name") else str(op)
+                if op_name in control_flow_ops:
+                    return instructions, list(constants) if constants is not None else []
+
+        self._consts = list(constants) if constants is not None else []
+        self._const_indexed = constants is not None
+
+        normalized: List[Instruction] = []
+        for instr in instructions:
+            if instr is None:
+                continue
+            if isinstance(instr, tuple) and len(instr) >= 2:
+                op = instr[0]
+                operand = instr[1] if len(instr) == 2 else tuple(instr[1:])
+                op_name = op.name if hasattr(op, "name") else op
+                normalized.append(Instruction(str(op_name), operand, 0))
+
+        optimized = self.optimize(normalized)
+
+        out: List[Tuple[Any, Any]] = []
+        for instr in optimized:
+            if instr.opcode == "LOAD_CONST" and self._const_indexed:
+                out.append((instr.opcode, instr.arg))
+            else:
+                out.append((instr.opcode, instr.arg))
+
+        return out, list(self._consts)
     
     def _apply_pattern(
         self,
@@ -199,9 +258,9 @@ class PeepholeOptimizer:
                            'BINARY_ADD', 'BINARY_SUB', 'BINARY_MUL',
                            'BINARY_DIV', 'BINARY_MOD')):
             
-            # Both operands must be constants
-            a = inst1.arg
-            b = inst2.arg
+            # Resolve constant values (supports constant pool indices)
+            a = self._resolve_const(inst1.arg)
+            b = self._resolve_const(inst2.arg)
             
             # Only fold numeric constants
             if not (isinstance(a, (int, float)) and isinstance(b, (int, float))):
@@ -228,7 +287,11 @@ class PeepholeOptimizer:
                     return False, [], 0
                 
                 # Replace with single LOAD_CONST
-                replacement = [Instruction('LOAD_CONST', result, inst1.lineno)]
+                if self._const_indexed:
+                    const_index = self._add_const(result)
+                    replacement = [Instruction('LOAD_CONST', const_index, inst1.lineno)]
+                else:
+                    replacement = [Instruction('LOAD_CONST', result, inst1.lineno)]
                 self.stats.constant_folds += 1
                 return True, replacement, 3
                 
@@ -237,6 +300,21 @@ class PeepholeOptimizer:
                 return False, [], 0
         
         return False, [], 0
+
+    def _resolve_const(self, arg: Any) -> Any:
+        if self._const_indexed and isinstance(arg, int) and self._consts is not None:
+            if 0 <= arg < len(self._consts):
+                return self._consts[arg]
+        return arg
+
+    def _add_const(self, value: Any) -> int:
+        if self._consts is None:
+            self._consts = []
+        for i, const in enumerate(self._consts):
+            if const == value and type(const) == type(value):
+                return i
+        self._consts.append(value)
+        return len(self._consts) - 1
     
     # ========== Dead Code Elimination ==========