"""Tests for code generation. Tests verify: - token-migration.AC7.1: Token stream mode emits IRAMWriteToken (not LoadInstToken) - token-migration.AC7.2: Token stream mode does not emit RouteSetToken - token-migration.AC7.3: Direct mode (PEConfig/SMConfig) still works Also tests original codegen AC8 criteria: - or1-asm.AC8.1: Direct mode produces valid PEConfig with correct IRAM contents - or1-asm.AC8.2: Direct mode produces valid SMConfig with initial cell values - or1-asm.AC8.3: Direct mode produces seed MonadTokens for const nodes with no incoming edges - or1-asm.AC8.4: Direct mode PEConfig includes route restrictions matching edge analysis - or1-asm.AC8.9: Program with no data_defs produces empty SM init section """ from asm.codegen import generate_direct, generate_tokens, AssemblyResult from asm.ir import ( IRGraph, IRNode, IREdge, IRDataDef, SystemConfig, SourceLoc, ResolvedDest, ) from cm_inst import ALUInst, Addr, ArithOp, MemOp, Port, RoutingOp, SMInst from tokens import IRAMWriteToken, MonadToken, SMToken from emu.types import PEConfig, SMConfig from sm_mod import Presence class TestTokenMigration: """Token migration acceptance criteria (AC7.1, AC7.2, AC7.3).""" def test_ac71_iram_write_token_in_stream(self): """AC7.1: Token stream mode emits IRAMWriteToken (not LoadInstToken). Tests that: - generate_tokens() produces IRAMWriteToken instances - At least one IRAMWriteToken is present for each PE with instructions """ node = IRNode( name="&add", opcode=ArithOp.ADD, pe=0, iram_offset=0, ctx=0, loc=SourceLoc(1, 1), ) system = SystemConfig(pe_count=1, sm_count=1) graph = IRGraph({"&add": node}, system=system) tokens = generate_tokens(graph) # Find IRAMWriteToken instances iram_write_tokens = [ t for t in tokens if isinstance(t, IRAMWriteToken) ] assert len(iram_write_tokens) > 0, "Should emit at least one IRAMWriteToken (AC7.1)" for token in iram_write_tokens: assert isinstance(token.instructions, tuple), "IRAMWriteToken should have instructions tuple" assert len(token.instructions) > 0, "IRAMWriteToken instructions should not be empty" def test_ac72_no_route_set_token(self): """AC7.2: Token stream mode does not emit RouteSetToken. Tests that: - generate_tokens() produces no RouteSetToken instances """ # Multi-PE graph to test routing node_pe0 = IRNode( name="&a", opcode=ArithOp.ADD, pe=0, iram_offset=0, ctx=0, dest_l=ResolvedDest( name="&b", addr=Addr(a=0, port=Port.L, pe=1), ), loc=SourceLoc(1, 1), ) node_pe1 = IRNode( name="&b", opcode=ArithOp.SUB, pe=1, iram_offset=0, ctx=0, loc=SourceLoc(2, 1), ) edge = IREdge(source="&a", dest="&b", port=Port.L, loc=SourceLoc(1, 1)) system = SystemConfig(pe_count=2, sm_count=1) graph = IRGraph( {"&a": node_pe0, "&b": node_pe1}, edges=[edge], system=system, ) tokens = generate_tokens(graph) # Verify no RouteSetToken (check by class name to be robust) route_set_tokens = [ t for t in tokens if type(t).__name__ == 'RouteSetToken' ] assert len(route_set_tokens) == 0, "Should not emit RouteSetToken (AC7.2)" def test_ac73_direct_mode_still_works(self): """AC7.3: Direct mode (PEConfig/SMConfig) still works. Tests that: - generate_direct() produces valid PEConfig with correct IRAM, route restrictions - generate_direct() produces valid SMConfig with initial cell values - Seed tokens are generated correctly """ data_def = IRDataDef( name="@val", sm_id=0, cell_addr=5, value=42, loc=SourceLoc(1, 1), ) node1 = IRNode( name="&a", opcode=ArithOp.ADD, pe=0, iram_offset=0, ctx=0, loc=SourceLoc(1, 1), ) node2 = IRNode( name="&b", opcode=RoutingOp.CONST, pe=0, iram_offset=1, ctx=0, const=99, loc=SourceLoc(2, 1), ) system = SystemConfig(pe_count=1, sm_count=1) graph = IRGraph( {"&a": node1, "&b": node2}, data_defs=[data_def], system=system, ) result = generate_direct(graph) # Verify PEConfig assert len(result.pe_configs) == 1 pe_config = result.pe_configs[0] assert pe_config.pe_id == 0 assert len(pe_config.iram) == 2 assert pe_config.allowed_pe_routes == {0} assert pe_config.allowed_sm_routes == set() # Verify SMConfig assert len(result.sm_configs) == 1 sm_config = result.sm_configs[0] assert sm_config.sm_id == 0 assert 5 in sm_config.initial_cells pres, val = sm_config.initial_cells[5] assert pres == Presence.FULL assert val == 42 # Verify seed tokens assert len(result.seed_tokens) == 1 seed = result.seed_tokens[0] assert isinstance(seed, MonadToken) assert seed.data == 99 class TestDirectMode: """AC8.1, AC8.2, AC8.3, AC8.4: Direct mode code generation.""" def test_ac81_simple_alu_instructions(self): """AC8.1: Two ALU nodes on PE0 produce PEConfig with correct IRAM. Tests that: - ALU instructions are correctly converted to ALUInst - They are placed in IRAM at assigned offsets """ # Create two simple dyadic ALU nodes add_node = IRNode( name="&add", opcode=ArithOp.ADD, pe=0, iram_offset=0, ctx=0, loc=SourceLoc(1, 1), ) sub_node = IRNode( name="&sub", opcode=ArithOp.SUB, pe=0, iram_offset=1, ctx=0, loc=SourceLoc(2, 1), ) system = SystemConfig(pe_count=1, sm_count=1) graph = IRGraph( {"&add": add_node, "&sub": sub_node}, system=system, ) result = generate_direct(graph) assert len(result.pe_configs) == 1 pe_config = result.pe_configs[0] assert pe_config.pe_id == 0 assert len(pe_config.iram) == 2 assert 0 in pe_config.iram assert 1 in pe_config.iram # Check the instruction types inst_0 = pe_config.iram[0] inst_1 = pe_config.iram[1] assert isinstance(inst_0, ALUInst) # Is an ALUInst assert isinstance(inst_1, ALUInst) # Is an ALUInst assert inst_0.op == ArithOp.ADD assert inst_1.op == ArithOp.SUB def test_ac82_data_defs_to_smconfig(self): """AC8.2: Data definitions produce SMConfig with initial cell values. Tests that: - Data defs with SM placement are converted to SMConfig - initial_cells dict contains correct (Presence.FULL, value) tuples """ data_def = IRDataDef( name="@val", sm_id=0, cell_addr=5, value=42, loc=SourceLoc(1, 1), ) graph = IRGraph({}, data_defs=[data_def], system=SystemConfig(1, 1)) result = generate_direct(graph) assert len(result.sm_configs) == 1 sm_config = result.sm_configs[0] assert sm_config.sm_id == 0 assert sm_config.initial_cells is not None assert 5 in sm_config.initial_cells pres, val = sm_config.initial_cells[5] assert pres == Presence.FULL assert val == 42 def test_ac83_const_node_seed_token(self): """AC8.3: CONST node with no incoming edges produces seed MonadToken. Tests that: - CONST nodes are detected - Nodes with no incoming edges are marked as seeds - MonadToken has correct target PE, offset, ctx, data """ const_node = IRNode( name="&seed", opcode=RoutingOp.CONST, pe=0, iram_offset=2, ctx=0, const=99, loc=SourceLoc(1, 1), ) graph = IRGraph({"&seed": const_node}, system=SystemConfig(1, 1)) result = generate_direct(graph) assert len(result.seed_tokens) == 1 token = result.seed_tokens[0] assert isinstance(token, MonadToken) assert token.target == 0 assert token.offset == 2 assert token.ctx == 0 assert token.data == 99 assert token.inline == False def test_ac84_route_restrictions(self): """AC8.4: Cross-PE edges produce correct allowed_pe_routes. Tests that: - Edges from PE0 to PE1 add PE1 to PE0's allowed_pe_routes - Self-routes are always included """ # PE0 node connecting to PE1 node node_pe0 = IRNode( name="&a", opcode=ArithOp.ADD, pe=0, iram_offset=0, ctx=0, dest_l=ResolvedDest( name="&b", addr=Addr(a=0, port=Port.L, pe=1), ), loc=SourceLoc(1, 1), ) node_pe1 = IRNode( name="&b", opcode=ArithOp.ADD, pe=1, iram_offset=0, ctx=0, loc=SourceLoc(2, 1), ) edge = IREdge(source="&a", dest="&b", port=Port.L, loc=SourceLoc(1, 1)) system = SystemConfig(pe_count=2, sm_count=1) graph = IRGraph( {"&a": node_pe0, "&b": node_pe1}, edges=[edge], system=system, ) result = generate_direct(graph) assert len(result.pe_configs) == 2 pe0_config = next(c for c in result.pe_configs if c.pe_id == 0) pe1_config = next(c for c in result.pe_configs if c.pe_id == 1) # PE0 should have routes to {0, 1} assert 0 in pe0_config.allowed_pe_routes assert 1 in pe0_config.allowed_pe_routes # PE1 should have route to {1} (self only, no incoming cross-PE edges) assert 1 in pe1_config.allowed_pe_routes def test_sm_instructions_in_iram(self): """Verify SMInst objects are correctly created and placed in IRAM. Tests that MemOp instructions produce SMInst in IRAM. """ sm_node = IRNode( name="&read", opcode=MemOp.READ, pe=0, iram_offset=0, ctx=0, sm_id=0, const=42, dest_l=ResolvedDest( name="&out", addr=Addr(a=1, port=Port.L, pe=0), ), loc=SourceLoc(1, 1), ) graph = IRGraph({"&read": sm_node}, system=SystemConfig(1, 1)) result = generate_direct(graph) assert len(result.pe_configs) == 1 pe_config = result.pe_configs[0] assert 0 in pe_config.iram inst = pe_config.iram[0] assert isinstance(inst, SMInst) # Is an SMInst assert inst.op == MemOp.READ assert inst.sm_id == 0 assert inst.const == 42 class TestTokenStream: """AC7.1, AC7.2, AC7.3: Token stream generation and ordering.""" def test_ac85_ac86_ac87_token_ordering(self): """AC7.1-7.2: Token stream emits SM init, IRAM writes, then seeds (no ROUTE_SET or LOAD_INST). Tests that: - SM init tokens come first - IRAM write tokens come next (IRAMWriteToken, not LoadInstToken) - Seed tokens come last - No RouteSetToken is present """ # Create a multi-PE graph with data_defs data_def = IRDataDef( name="@val", sm_id=0, cell_addr=5, value=42, loc=SourceLoc(1, 1), ) node1 = IRNode( name="&a", opcode=ArithOp.ADD, pe=0, iram_offset=0, ctx=0, loc=SourceLoc(1, 1), ) node2 = IRNode( name="&b", opcode=RoutingOp.CONST, pe=0, iram_offset=1, ctx=0, const=10, loc=SourceLoc(2, 1), ) system = SystemConfig(pe_count=1, sm_count=1) graph = IRGraph( {"&a": node1, "&b": node2}, data_defs=[data_def], system=system, ) tokens = generate_tokens(graph) # Find positions of token types smtoken_indices = [ i for i, t in enumerate(tokens) if isinstance(t, SMToken) ] iram_write_indices = [ i for i, t in enumerate(tokens) if isinstance(t, IRAMWriteToken) ] seed_indices = [ i for i, t in enumerate(tokens) if isinstance(t, MonadToken) ] # Verify order: SM < IRAM write < seed assert smtoken_indices, "Should have at least one SM token" assert iram_write_indices, "Should have at least one IRAM write token" assert seed_indices, "Should have at least one seed token" assert max(smtoken_indices) < min(iram_write_indices), "SM tokens should come before IRAM write tokens" assert max(iram_write_indices) < min(seed_indices), "IRAM write tokens should come before seed tokens" def test_ac88_tokens_are_valid(self): """AC7.3: Generated tokens in direct mode are valid and direct mode PEConfig/SMConfig still works. Tests that: - All tokens have required fields set - Token structure matches emulator expectations - Direct mode produces valid PEConfig/SMConfig - Tokens can be injected into an emulator System and execution completes """ from emu.network import build_topology import simpy data_def = IRDataDef( name="@val", sm_id=0, cell_addr=5, value=42, loc=SourceLoc(1, 1), ) node = IRNode( name="&add", opcode=ArithOp.ADD, pe=0, iram_offset=0, ctx=0, loc=SourceLoc(1, 1), ) system = SystemConfig(pe_count=1, sm_count=1) graph = IRGraph( {"&add": node}, data_defs=[data_def], system=system, ) result = generate_direct(graph) tokens = generate_tokens(graph) # Verify direct mode result structure assert isinstance(result, AssemblyResult) assert len(result.pe_configs) == 1 assert result.pe_configs[0].pe_id == 0 assert len(result.sm_configs) == 1 assert result.sm_configs[0].sm_id == 0 # Build emulator system from AssemblyResult configs env = simpy.Environment() emu_system = build_topology( env, result.pe_configs, result.sm_configs, fifo_capacity=16, ) # Inject tokens into the system following the new sequence: # 1. SM init tokens # 2. IRAM write tokens # 3. Seed MonadTokens for token in tokens: if isinstance(token, SMToken): emu_system.inject(token) elif isinstance(token, IRAMWriteToken): emu_system.inject(token) elif isinstance(token, MonadToken): emu_system.inject(token) # Run the simulation for enough steps to complete initialization env.run(until=1000) # Verify token structure for token in tokens: if isinstance(token, SMToken): assert isinstance(token.target, int) assert isinstance(token.addr, int) assert isinstance(token.op, MemOp) assert token.op == MemOp.WRITE elif isinstance(token, IRAMWriteToken): assert isinstance(token.target, int) assert isinstance(token.offset, int) assert isinstance(token.ctx, int) assert isinstance(token.data, int) assert isinstance(token.instructions, tuple) elif isinstance(token, MonadToken): assert isinstance(token.target, int) assert isinstance(token.offset, int) assert isinstance(token.ctx, int) assert isinstance(token.data, int) class TestEdgeCases: """AC8.9, AC8.10: Edge cases for code generation.""" def test_ac89_no_data_defs(self): """AC8.9: Program with no data_defs produces no SMConfig or SM tokens. Tests that: - sm_configs list is empty - Token stream has no SMTokens """ node = IRNode( name="&add", opcode=ArithOp.ADD, pe=0, iram_offset=0, ctx=0, loc=SourceLoc(1, 1), ) system = SystemConfig(pe_count=1, sm_count=1) graph = IRGraph({"&add": node}, system=system) result = generate_direct(graph) assert len(result.sm_configs) == 0 tokens = generate_tokens(graph) sm_tokens = [t for t in tokens if isinstance(t, SMToken)] assert len(sm_tokens) == 0 def test_ac810_single_pe_self_route(self): """AC7.2: Single-PE program produces IRAM writes with no RouteSetToken. Tests that: - allowed_pe_routes contains only the PE's own ID (in direct mode) - Token stream has no RouteSetToken (route restrictions are not emitted) - Token stream has IRAMWriteToken instead """ node = IRNode( name="&add", opcode=ArithOp.ADD, pe=0, iram_offset=0, ctx=0, loc=SourceLoc(1, 1), ) system = SystemConfig(pe_count=1, sm_count=1) graph = IRGraph({"&add": node}, system=system) result = generate_direct(graph) assert len(result.pe_configs) == 1 pe_config = result.pe_configs[0] assert pe_config.allowed_pe_routes == {0} tokens = generate_tokens(graph) # Verify no RouteSetToken (AC7.2) route_set_tokens = [ t for t in tokens if type(t).__name__ == 'RouteSetToken' # Check by class name to avoid import ] assert len(route_set_tokens) == 0, "RouteSetToken should not be in token stream (AC7.2)" # Verify IRAMWriteToken is present (AC7.1) iram_write_tokens = [ t for t in tokens if isinstance(t, IRAMWriteToken) ] assert len(iram_write_tokens) == 1, "Should have exactly one IRAMWriteToken per PE" token = iram_write_tokens[0] assert token.target == 0 def test_multiple_data_defs_same_sm(self): """Multiple data_defs targeting same SM produce single SMConfig. Tests that: - Multiple data_defs for SM0 are merged into single SMConfig - initial_cells contains all entries """ data_def1 = IRDataDef( name="@val1", sm_id=0, cell_addr=5, value=42, loc=SourceLoc(1, 1), ) data_def2 = IRDataDef( name="@val2", sm_id=0, cell_addr=10, value=99, loc=SourceLoc(2, 1), ) system = SystemConfig(pe_count=1, sm_count=1) graph = IRGraph({}, data_defs=[data_def1, data_def2], system=system) result = generate_direct(graph) assert len(result.sm_configs) == 1 sm_config = result.sm_configs[0] assert sm_config.sm_id == 0 assert len(sm_config.initial_cells) == 2 assert sm_config.initial_cells[5] == (Presence.FULL, 42) assert sm_config.initial_cells[10] == (Presence.FULL, 99) def test_const_node_with_incoming_edge_not_seed(self): """CONST node with incoming edge is not a seed token. Tests that: - Only CONST nodes with NO incoming edges produce seed_tokens """ source_node = IRNode( name="&src", opcode=ArithOp.ADD, pe=0, iram_offset=0, ctx=0, loc=SourceLoc(1, 1), ) const_node = IRNode( name="&const", opcode=RoutingOp.CONST, pe=0, iram_offset=1, ctx=0, const=5, loc=SourceLoc(2, 1), ) edge = IREdge(source="&src", dest="&const", port=Port.L, loc=SourceLoc(1, 1)) system = SystemConfig(pe_count=1, sm_count=1) graph = IRGraph( {"&src": source_node, "&const": const_node}, edges=[edge], system=system, ) result = generate_direct(graph) # The CONST node has an incoming edge, so it should NOT be a seed assert len(result.seed_tokens) == 0 class TestMultiPERouting: """Extended tests for multi-PE routing scenarios.""" def test_multi_pe_route_computation(self): """Multi-PE graph with multiple cross-PE edges. Tests route computation across multiple PEs with various edge patterns. """ # Create a 3-PE system with cross-PE edges node_pe0 = IRNode( name="&a", opcode=ArithOp.ADD, pe=0, iram_offset=0, ctx=0, dest_l=ResolvedDest( name="&b", addr=Addr(a=0, port=Port.L, pe=1), ), loc=SourceLoc(1, 1), ) node_pe1 = IRNode( name="&b", opcode=ArithOp.SUB, pe=1, iram_offset=0, ctx=0, dest_l=ResolvedDest( name="&c", addr=Addr(a=0, port=Port.L, pe=2), ), loc=SourceLoc(2, 1), ) node_pe2 = IRNode( name="&c", opcode=ArithOp.INC, pe=2, iram_offset=0, ctx=0, loc=SourceLoc(3, 1), ) edge1 = IREdge(source="&a", dest="&b", port=Port.L, loc=SourceLoc(1, 1)) edge2 = IREdge(source="&b", dest="&c", port=Port.L, loc=SourceLoc(2, 1)) system = SystemConfig(pe_count=3, sm_count=1) graph = IRGraph( {"&a": node_pe0, "&b": node_pe1, "&c": node_pe2}, edges=[edge1, edge2], system=system, ) result = generate_direct(graph) pe0_config = next(c for c in result.pe_configs if c.pe_id == 0) pe1_config = next(c for c in result.pe_configs if c.pe_id == 1) pe2_config = next(c for c in result.pe_configs if c.pe_id == 2) # PE0 -> PE1 assert 1 in pe0_config.allowed_pe_routes # PE1 -> PE2 assert 2 in pe1_config.allowed_pe_routes # PE2 has no outgoing edges assert pe2_config.allowed_pe_routes == {2}