Implemented SourceFile and Solution in the rest of the program

esbmc_ai/__main__.py

@@ -3,24 +3,23 @@
 # Author: Yiannis Charalambous 2023
 
 import os
+from pathlib import Path
 import re
 import sys
 
 # Enables arrow key functionality for input(). Do not remove import.
 import readline
+from typing import Optional
+
+_ = readline
 
 import argparse
 from langchain.base_language import BaseLanguageModel
 
 
 import esbmc_ai.config as config
 from esbmc_ai import __author__, __version__
-from esbmc_ai.frontend.solution import (
-    SourceFile,
-    get_main_source_file,
-    set_main_source_file,
-    get_main_source_file_path,
-)
+from esbmc_ai.solution import SourceFile, Solution, get_solution, init_solution
 
 from esbmc_ai.commands import (
     ChatCommand,
@@ -123,7 +122,34 @@ def init_commands_list() -> None:
 
 
 def update_solution(source_code: str) -> None:
-    set_main_source_file(SourceFile(get_main_source_file_path(), source_code))
+    get_solution().files[0].update_content(content=source_code, reset_changes=True)
+
+
+def _run_esbmc(source_file: SourceFile, anim: Optional[LoadingWidget] = None) -> str:
+    assert source_file.file_path
+
+    if anim:
+        anim.start("ESBMC is processing... Please Wait")
+    exit_code, esbmc_output = esbmc(
+        path=source_file.file_path,
+        esbmc_params=config.esbmc_params,
+        timeout=config.verifier_timeout,
+    )
+    if anim:
+        anim.stop()
+
+    # ESBMC will output 0 for verification success and 1 for verification
+    # failed, if anything else gets thrown, it's an ESBMC error.
+    if not config.allow_successful and exit_code == 0:
+        printv("Success!")
+        print(esbmc_output)
+        sys.exit(0)
+    elif exit_code != 0 and exit_code != 1:
+        printv(f"ESBMC exit code: {exit_code}")
+        printv(f"ESBMC Output:\n\n{esbmc_output}")
+        sys.exit(1)
+
+    return esbmc_output
 
 
 def init_commands() -> None:
@@ -137,25 +163,34 @@ def init_commands() -> None:
     fix_code_command.on_solution_signal.add_listener(update_solution)
 
 
-def _run_command_mode(
-    command: ChatCommand,
-    args: list[str],
-    esbmc_output: str,
-    source_code: str,
-) -> None:
+def _run_command_mode(command: ChatCommand, args: argparse.Namespace) -> None:
+    path_arg: Path = Path(args.filename)
+
+    solution: Solution
+    if path_arg.is_dir():
+        solution = init_solution(
+            Solution(
+                [path for path in path_arg.glob("**/*") if path.is_file() and path.name]
+            )
+        )
+    else:
+        solution = init_solution(Solution([path_arg]))
+
     match command.command_name:
         case fix_code_command.command_name:
-            error, solution = fix_code_command.execute(
-                file_name=get_main_source_file_path(),
-                source_code=source_code,
-                esbmc_output=esbmc_output,
-            )
+            for source_file in solution.files:
+                # Run ESBMC first round
+                esbmc_output: str = _run_esbmc(source_file)
+                source_file.assign_verifier_output(esbmc_output)
 
-            if error:
-                print("Failed all attempts...")
-                sys.exit(1)
-            else:
-                print(solution)
+                error, source_file_solution = fix_code_command.execute(
+                    source_file=source_file
+                )
+
+                if error:
+                    print("Failed all attempts...")
+                else:
+                    print(source_file_solution)
         case _:
             command.execute()
     sys.exit(0)
@@ -249,7 +284,7 @@ def main() -> None:
         + "}",
     )
 
-    args = parser.parse_args()
+    args: argparse.Namespace = parser.parse_args()
 
     print(f"ESBMC-AI {__version__}")
     print(f"Made by {__author__}")
@@ -270,37 +305,7 @@ def main() -> None:
     printv("Reading source code...")
     print(f"Running ESBMC with {config.esbmc_params}\n")
 
-    # Cast to string (for language servers)
-    args.filename = str(args.filename)
-
-    # Read source code
-    with open(args.filename, mode="r") as file:
-        # Add the main source file to the solution explorer.
-        set_main_source_file(SourceFile(args.filename, file.read()))
-
-    anim.start("ESBMC is processing... Please Wait")
-    exit_code, esbmc_output = esbmc(
-        path=get_main_source_file_path(),
-        esbmc_params=config.esbmc_params,
-        timeout=config.verifier_timeout,
-    )
-    anim.stop()
-
-    # Print verbose lvl 2
-    print_horizontal_line(2)
-    printvv(esbmc_output)
-    print_horizontal_line(2)
-
-    # ESBMC will output 0 for verification success and 1 for verification
-    # failed, if anything else gets thrown, it's an ESBMC error.
-    if not config.allow_successful and exit_code == 0:
-        print("Success!")
-        print(esbmc_output)
-        sys.exit(0)
-    elif exit_code != 0 and exit_code != 1:
-        print(f"ESBMC exit code: {exit_code}")
-        print(f"ESBMC Output:\n\n{esbmc_output}")
-        sys.exit(1)
+    assert isinstance(args.filename, str)
 
     # Command mode: Check if command is called and call it.
     # If not, then continue to user mode.
@@ -310,14 +315,45 @@ def main() -> None:
             print("Running Command:", command)
             for idx, command_name in enumerate(command_names):
                 if command == command_name:
-                    _run_command_mode(
-                        command=commands[idx],
-                        args=[],  # NOTE: Currently not supported...
-                        source_code=get_main_source_file().content,
-                        esbmc_output=esbmc_output,
-                    )
+                    _run_command_mode(command=commands[idx], args=args)
             sys.exit(0)
 
+    # ===========================================
+    # User Mode (Supports only 1 file)
+    # ===========================================
+
+    # Init Solution
+    solution: Solution
+    # Determine if we are processing one file versus multiple files
+    path_arg: Path = Path(args.filename)
+    if path_arg.is_dir():
+        # Load only files.
+        print(
+            "Processing multiple files is not supported in User Mode."
+            "Call a command using -c to process directories"
+        )
+        sys.exit(1)
+    else:
+        # Add the main source file to the solution explorer.
+        solution = init_solution(Solution([path_arg]))
+    del path_arg
+
+    # Assert that we have one file and one filepath
+    assert len(solution.files) == 1
+
+    source_file: SourceFile = solution.files[0]
+    assert source_file.file_path
+
+    esbmc_output: str = _run_esbmc(source_file, anim)
+
+    # Print verbose lvl 2
+    print_horizontal_line(2)
+    printvv(esbmc_output)
+    print_horizontal_line(2)
+
+    source_file.assign_verifier_output(esbmc_output)
+    del esbmc_output
+
     printv(f"Initializing the LLM: {config.ai_model.name}\n")
     chat_llm: BaseLanguageModel = config.ai_model.create_llm(
         api_keys=config.api_keys,
@@ -332,8 +368,8 @@ def main() -> None:
         ai_model_agent=config.chat_prompt_user_mode,
         ai_model=config.ai_model,
         llm=chat_llm,
-        source_code=get_main_source_file().content,
-        esbmc_output=esbmc_output,
+        source_code=source_file.latest_content,
+        esbmc_output=source_file.latest_verifier_output,
         set_solution_messages=config.chat_prompt_user_mode.scenarios["set_solution"],
     )
 
@@ -345,7 +381,6 @@ def main() -> None:
     if len(config.chat_prompt_user_mode.initial_prompt) > 0:
         printv("Using initial prompt from file...\n")
         anim.start("Model is parsing ESBMC output... Please Wait")
-        # TODO Make protected
         response = chat.send_message(
             message=config.chat_prompt_user_mode.initial_prompt,
         )
@@ -375,17 +410,17 @@ def main() -> None:
                 print()
                 print("ESBMC-AI will generate a fix for the code...")
 
-                error, solution = fix_code_command.execute(
-                    file_name=get_main_source_file_path(),
-                    source_code=get_main_source_file().content,
-                    esbmc_output=esbmc_output,
+                error: bool
+                source_file_solution: str
+                error, source_file_solution = fix_code_command.execute(
+                    source_file=source_file
                 )
 
                 if not error:
                     print(
                         "\n\nESBMC-AI: Here is the corrected code, verified with ESBMC:"
                     )
-                    print(f"```\n{solution}\n```")
+                    print(f"```\n{source_file_solution}\n```")
                 continue
             else:
                 # Commands without parameters or returns are handled automatically.

esbmc_ai/chats/solution_generator.py

@@ -8,7 +8,7 @@
 
 from esbmc_ai.chat_response import ChatResponse, FinishReason
 from esbmc_ai.config import ChatPromptSettings, DynamicAIModelAgent
-from esbmc_ai.frontend.solution import apply_line_patch
+from esbmc_ai.solution import SourceFile
 
 from esbmc_ai.ai_models import AIModel
 from .base_chat_interface import BaseChatInterface
@@ -205,6 +205,8 @@ def generate_solution(self) -> tuple[str, FinishReason]:
                 assert (
                     line
                 ), "fix code command: error line could not be found to apply brutal patch replacement"
-                solution = apply_line_patch(self.source_code_raw, solution, line, line)
+                solution = SourceFile.apply_line_patch(
+                    self.source_code_raw, solution, line, line
+                )
 
         return solution, response.finish_reason

esbmc_ai/commands/fix_code_command.py

@@ -8,6 +8,7 @@
 from esbmc_ai.chats import LatestStateSolutionGenerator, SolutionGenerator
 from esbmc_ai.chats.solution_generator import ESBMCTimedOutException
 from esbmc_ai.reverse_order_solution_generator import ReverseOrderSolutionGenerator
+from esbmc_ai.solution import SourceFile
 
 from .chat_command import ChatCommand
 from .. import config
@@ -41,12 +42,11 @@ def print_raw_conversation() -> None:
             print("\n" + "\n\n".join(messages))
             print("ESBMC-AI Notice: End of raw conversation")
 
-        file_name: str = kwargs["file_name"]
-        source_code: str = kwargs["source_code"]
-        esbmc_output: str = kwargs["esbmc_output"]
+        source_file: SourceFile = kwargs["source_file"]
+        assert source_file.file_path
 
         # Parse the esbmc output here and determine what "Scenario" to use.
-        scenario: str = esbmc_get_error_type(esbmc_output)
+        scenario: str = esbmc_get_error_type(source_file.initial_verifier_output)
 
         printv(f"Scenario: {scenario}")
         printv(
@@ -105,8 +105,8 @@ def print_raw_conversation() -> None:
 
         try:
             solution_generator.update_state(
-                source_code=source_code,
-                esbmc_output=esbmc_output,
+                source_code=source_file.latest_content,
+                esbmc_output=source_file.latest_verifier_output,
             )
         except ESBMCTimedOutException:
             print("error: ESBMC has timed out...")
@@ -126,48 +126,53 @@ def print_raw_conversation() -> None:
                 if finish_reason == FinishReason.length:
                     solution_generator.compress_message_stack()
                 else:
-                    source_code = llm_solution
+                    source_file.update_content(llm_solution)
                     break
 
             # Print verbose lvl 2
             printvv("\nESBMC-AI Notice: Source Code Generation:")
             print_horizontal_line(2)
-            printvv(source_code)
+            printvv(source_file.latest_content)
             print_horizontal_line(2)
             printvv("")
 
             # Pass to ESBMC, a workaround is used where the file is saved
             # to a temporary location since ESBMC needs it in file format.
             self.anim.start("Verifying with ESBMC... Please Wait")
             exit_code, esbmc_output = esbmc_load_source_code(
-                file_path=file_name,
-                source_code=source_code,
+                file_path=source_file.file_path,
+                source_code=source_file.latest_content,
                 esbmc_params=config.esbmc_params,
                 auto_clean=config.temp_auto_clean,
                 timeout=config.verifier_timeout,
             )
             self.anim.stop()
 
+            source_file.assign_verifier_output(esbmc_output)
+            del esbmc_output
+
             # Print verbose lvl 2
             printvv("\nESBMC-AI Notice: ESBMC Output:")
             print_horizontal_line(2)
-            printvv(esbmc_output)
+            printvv(source_file.latest_verifier_output)
             print_horizontal_line(2)
 
             # Solution found
             if exit_code == 0:
-                self.on_solution_signal.emit(source_code)
+                self.on_solution_signal.emit(source_file.latest_content)
 
                 if config.raw_conversation:
                     print_raw_conversation()
 
                 printv("ESBMC-AI Notice: Successfully verified code")
 
-                return False, source_code
+                return False, source_file.latest_content
 
             try:
                 # Update state
-                solution_generator.update_state(source_code, esbmc_output)
+                solution_generator.update_state(
+                    source_file.latest_content, source_file.latest_verifier_output
+                )
             except ESBMCTimedOutException:
                 if config.raw_conversation:
                     print_raw_conversation()

tests/test_solution.py

@@ -1,13 +1,16 @@
 # Author: Yiannis Charalambous
 
-from esbmc_ai.frontend.solution import apply_line_patch
+from esbmc_ai.solution import SourceFile
 
 
 def test_apply_line_patch() -> None:
     text = "\n".join(["a", "b", "c", "d", "e", "f", "g"])
     answer = "\n".join(["a", "b", "1", "g"])
-    assert apply_line_patch(text, "1", 2, 5) == answer
+    assert SourceFile.apply_line_patch(text, "1", 2, 5) == answer
 
     text = "\n".join(["a", "b", "c", "d", "e", "f", "g"])
     answer = "\n".join(["a", "b", "c", "1", "e", "f", "g"])
-    assert apply_line_patch(text, "1", 3, 3) == answer
+    assert SourceFile.apply_line_patch(text, "1", 3, 3) == answer
+
+
+# TODO Add more tests for solution and sourcefile