File size: 16,067 Bytes
08aa12c
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.

"""

Rtl Debugger Environment Implementation.



On reset(), a random task is loaded from the tasks/ directory.

Each task folder must contain:

    design_buggy.v  - the immutable original buggy module shown to the agent

    makefile        - the cocotb makefile (must reference design_active.v)

    test_*.py       - the cocotb testbench(es) that produce result.json



design_active.v is created on reset() as a mutable working copy of design_buggy.v.

The agent overwrites design_active.v each step, so it iterates on its own previous fix.

"""

import difflib
import json
import os
import random
import re
import shutil
import subprocess
import traceback
from uuid import uuid4

from openenv.core.env_server.interfaces import Environment
from openenv.core.env_server.types import State

try:
    from ..models import EditOp, RtlDebuggerAction, RtlDebuggerObservation
except ImportError:
    from models import EditOp, RtlDebuggerAction, RtlDebuggerObservation

from .graders import get_grader

# Tasks directory is usually at the project root.
# We first try the current working directory (e.g. /app/env in Docker), 
# then fall back to the package-relative path.
_TASKS_DIR = os.path.join(os.getcwd(), "tasks")
if not os.path.isdir(_TASKS_DIR):
    _TASKS_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "tasks"))


def _list_tasks() -> list[str]:
    """Return sorted list of task directory paths found under _TASKS_DIR.



    A valid task must contain design_buggy.v and a makefile.

    """
    if not os.path.isdir(_TASKS_DIR):
        return []
    return sorted(
        os.path.join(_TASKS_DIR, name)
        for name in os.listdir(_TASKS_DIR)
        if os.path.isdir(os.path.join(_TASKS_DIR, name))
        and os.path.exists(os.path.join(_TASKS_DIR, name, "design_buggy.v"))
        and os.path.exists(os.path.join(_TASKS_DIR, name, "makefile"))
    )


def _read(path: str) -> str:
    with open(path, "r") as f:
        return f.read()


def _levenshtein_line_distance(original: str, modified: str) -> int:
    """

    Compute a line-based edit distance between two strings.

    Treats line addition, deletion, and substitution as 1 operation each.

    """
    import difflib
    orig_lines = original.splitlines()
    mod_lines = modified.splitlines()
    
    matcher = difflib.SequenceMatcher(None, orig_lines, mod_lines)
    distance = 0
    for tag, i1, i2, j1, j2 in matcher.get_opcodes():
        if tag == 'replace':
            # substitution of a block of lines
            distance += max(i2 - i1, j2 - j1)
        elif tag == 'insert':
            distance += (j2 - j1)
        elif tag == 'delete':
            distance += (i2 - i1)
            
    return distance


def _numbered_code(file_path: str) -> str:
    """Return the contents of a file with 1-indexed line numbers prepended."""
    lines = _read(file_path).splitlines()
    return "\n".join(f"{i + 1:4d}| {line}" for i, line in enumerate(lines))


def _apply_edits(file_path: str, edits: list[EditOp]) -> str:
    """Apply a list of EditOp operations to a file and write back.



    Line numbers in each EditOp refer to the file state BEFORE any edits

    in this batch.  Operations are applied in reverse line-number order so

    that earlier indices stay valid as later lines are modified.



    Supports multi-line edits via ``end_line`` (range) and

    newlines inside ``new_content``.



    Returns the patched file content (also written to *file_path*).

    """
    lines = _read(file_path).splitlines()

    # Sort by line_number descending so later edits don't shift earlier indices.
    sorted_edits = sorted(edits, key=lambda e: e.line_number, reverse=True)

    for edit in sorted_edits:
        start = edit.line_number - 1          # 0-indexed start
        end = (edit.end_line or edit.line_number) - 1  # 0-indexed end (inclusive)
        end = min(end, len(lines) - 1)        # clamp to file length

        if edit.op == "replace":
            if 0 <= start < len(lines):
                new_lines = edit.new_content.split("\n")
                lines[start:end + 1] = new_lines
        elif edit.op == "insert_after":
            insert_pos = max(start + 1, 0) if edit.line_number > 0 else 0
            new_lines = edit.new_content.split("\n")
            for i, nl in enumerate(new_lines):
                lines.insert(insert_pos + i, nl)
        elif edit.op == "delete":
            if 0 <= start < len(lines):
                del lines[start:end + 1]

    patched = "\n".join(lines) + "\n"
    with open(file_path, "w") as f:
        f.write(patched)
    return patched


class RtlDebuggerEnvironment(Environment):
    """

    RTL Debugging environment.



    On reset(), a random task is selected from the tasks/ directory.

    - design_buggy.v is read and shown to agent as the original reference.

    - design_active.v is created as a fresh mutable copy of design_buggy.v.



    On each step(), the agent submits revised code which:

      1. Is written to design_active.v (overwrites the previous attempt).

      2. Is compiled with iverilog — fatal gate if it fails.

      3. Is simulated via `make` (cocotb + icarus).

      4. Has result.json parsed for pass/fail metrics.

      5. Levenshtein distance is measured from the *previous* design_active.v.



    A separate grader() function computes the final normalized score [0,1].

    """

    SUPPORTS_CONCURRENT_SESSIONS: bool = True

    def __init__(self):
        """Initialize the rtl_debugger environment."""
        self._state = State(episode_id=str(uuid4()), step_count=0)
        self._task_dir: str = ""

        self._task_context: str = ""  # contents of context.md
        self._task_name: str = ""

    @property
    def _active_design_path(self) -> str:
        return os.path.join(self._task_dir, "design_active.v")

    def reset(self, options: dict | None = None) -> RtlDebuggerObservation:  # type: ignore[override]
        """

        Reset the environment and load a task.

        Runs the initial 'buggy' design through simulation to provide baseline feedback.

        """
        self._state = State(episode_id=str(uuid4()), step_count=0)

        tasks = _list_tasks()
        if not tasks:
            return RtlDebuggerObservation(
                task_id="",

                feedback=f"No tasks found in {_TASKS_DIR}.",
                compiled=False,
                passed_tests=False,
                done=False,
                reward=0.0,
            )

        # Handle explicit task selection via options
        selected_task_dir = None
        if options and "TASK_NAME" in options:
            target_name = options["TASK_NAME"]
            for t_dir in tasks:
                if os.path.basename(t_dir) == target_name:
                    selected_task_dir = t_dir
                    break
        
        if not selected_task_dir:
            selected_task_dir = random.choice(tasks)

        self._task_dir = selected_task_dir
        self._task_name = os.path.basename(self._task_dir)
        buggy_path = os.path.join(self._task_dir, "design_buggy.v")

        
        context_path = os.path.join(self._task_dir, "context.md")
        if os.path.exists(context_path):
            self._task_context = _read(context_path)
        else:
            self._task_context = ""

        # Initialise design_active.v from the buggy design
        shutil.copy(buggy_path, self._active_design_path)

        result_path = os.path.join(self._task_dir, "result.json")
        if os.path.exists(result_path):
            os.remove(result_path)

        # Run baseline simulation to provide initial feedback
        baseline_obs = self._run_evaluation(lev_distance=0)
        
        # Override baseline feedback to include a header
        baseline_obs.feedback = f"--- (Buggy code feedback) ---\n" + baseline_obs.feedback
        return baseline_obs

    def step(self, action: RtlDebuggerAction) -> RtlDebuggerObservation:  # type: ignore[override]
        """

        Apply the agent's line-edit operations, then compile and simulate.

        """
        self._state.step_count += 1

        try:
            # --- Snapshot previous state for Levenshtein ---
            prev_active = _read(self._active_design_path)

            # --- Apply edits to design_active.v ---
            patched = _apply_edits(self._active_design_path, action.edits)

            lev_distance = _levenshtein_line_distance(prev_active, patched)

            return self._run_evaluation(lev_distance=lev_distance)

        except Exception as exc:
            traceback.print_exc()
            return RtlDebuggerObservation(
                task_id=self._task_name,

                numbered_code=_numbered_code(self._active_design_path),
                feedback=f"[Internal Server Error] {type(exc).__name__}: {exc}",
                compiled=False,
                passed_tests=False,
                pass_rate=0.0,
                progress_ratio=0,
                reward=-1.0,
                done=False,
            )

    def _run_evaluation(self, lev_distance: int) -> RtlDebuggerObservation:
        """

        Internal helper to compile, simulate, and score the current design_active.v.

        Used by both reset() and step().

        """
        # --- 2. Run simulation via make ---
        # Note: 'make' handles both compilation and simulation via cocotb.
        try:
            make_res = subprocess.run(
                ["make"],
                cwd=self._task_dir,
                capture_output=True,
                text=True,
                timeout=60,
            )
            make_feedback = make_res.stdout + make_res.stderr
        except subprocess.TimeoutExpired:
            return RtlDebuggerObservation(
                task_id=self._task_name,

                numbered_code=_numbered_code(self._active_design_path),
                task_context=self._task_context,
                feedback="[Simulation Timeout] Your design likely contains an infinite loop. Simulation killed after 60s.",
                compiled=False,
                passed_tests=False,
                pass_rate=0.0,
                progress_ratio=0,
                levenshtein_distance=lev_distance,
                reward=-5.0,
                done=False,
            )

        # --- 3. Parse result.json (written by cocotb into the task dir) ---
        result_path = os.path.join(self._task_dir, "result.json")
        if not os.path.exists(result_path):
            error_lines = []
            for line in make_feedback.splitlines():
                lower_line = line.lower()
                if any(kw in lower_line for kw in ["error", "exception", "fatal", "traceback", "assert"]):
                    error_lines.append(line)
            
            # If no explicit error keyword found, just keep the tail of the log
            if not error_lines:
                error_lines = make_feedback.splitlines()[-15:]
            
            filtered_feedback = "\n".join(error_lines)


            return RtlDebuggerObservation(
                task_id=self._task_name,

                numbered_code=_numbered_code(self._active_design_path),
                task_context=self._task_context,
                feedback=f"=== Simulation/Build Error ===\nSimulation crashed before producing results.\nError Log:\n{filtered_feedback}",
                compiled=False,
                passed_tests=False,
                pass_rate=0.0,
                progress_ratio=0,
                levenshtein_distance=lev_distance,
                reward=-3.0,
                done=False,
            )

        with open(result_path, "r") as f:
            result_data = json.load(f)

        num_tests = result_data.get("num_tests", 1)
        num_passed = result_data.get("num_passed", 0)
        pass_rate = num_passed / num_tests if num_tests > 0 else 0.0
        passed_all = result_data.get("passed", False)

        # Base reward for overall correctness
        reward = pass_rate * 12.0
        
        # Bonus reward for sequential correctness (fixing earliest bugs first)
        results_list = result_data.get("results", [])
        sequential_passes = 0
        for r in results_list:
            if r.get("pass", True):
                sequential_passes += 1
            else:
                break
                
        progress_ratio  = sequential_passes / max(num_tests, 1)
        reward += progress_ratio * 3.0  
        reward -= min(lev_distance * 0.01, 1.0) 

        if lev_distance == 0 and self._state.step_count > 0:
            reward -= 1.5
        if passed_all:
            reward += 20.0

        seq_count = result_data.get("sequence_correctness")
        trans_count = result_data.get("transition_correctness")

        if seq_count and trans_count:
            feedback_lines = [
                f"Compiled: OK | Tests: {num_passed}/{num_tests} passed | Seq: {seq_count}/{num_tests} | Trans: {trans_count}/{num_tests} | First failure after: {progress_ratio:.2f} of sequence"
            ]
        else:
            feedback_lines = [
                f"Compiled: OK | Tests: {num_passed}/{num_tests} passed | First failure after: {progress_ratio:.2f} of sequence"
            ]
        
        # Parse failed test cases from result.json
        if "results" in result_data:
            failed_tests = [r for r in result_data["results"] if not r.get("pass", True)]
            if failed_tests:
                feedback_lines.append("\n=== Failed Test Cases ===")
                for ft in failed_tests[:10]:
                    inp = ft.get("inputs") or ft.get("input", {})
                    inputs_str = ", ".join(f"{k}={v}" for k, v in inp.items()) if isinstance(inp, dict) else str(inp)
                    cycle = ft.get("test_id", "?")
                    expected = ft.get("expected_output", "?")
                    actual = ft.get("actual_output", "?")
                    
                    state_info = ""
                    if "transition_to" in ft:
                        was_s  = ft.get("actual_state_meaning", f"S{ft.get('actual_state', '?')}")
                        went_s = ft.get("transition_to_meaning", f"S{ft.get('transition_to', '?')}")
                        exp_s  = ft.get("expected_next_state_meaning", f"S{ft.get('expected_next_state', '?')}")
                        state_info = f" | was {was_s}, went to {went_s} (expected {exp_s})"
                    
                    feedback_lines.append(f"- Cycle {cycle}: [{inputs_str}] -> Expected: {expected}, Got: {actual}{state_info}")

        feedback = "\n".join(filter(None, feedback_lines))
        grader = get_grader(self._task_name)
        final_score = grader(self._state, self)

        print(feedback)

        return RtlDebuggerObservation(
            task_id=self._task_name,

            numbered_code=_numbered_code(self._active_design_path),
            task_context=self._task_context,
            feedback=feedback,
            compiled=True,
            passed_tests=passed_all,
            pass_rate=pass_rate,
            progress_ratio=progress_ratio,
            levenshtein_distance=lev_distance,
            reward=reward,
            done=passed_all,
            score=final_score,
        )

    @property
    def state(self) -> State:
        """

        Get the current environment state.



        Returns:

            Current State with episode_id and step_count

        """
        return self._state