assemble_tests.py

"""Script for assembling all of the test programs in /prog and placing them in /bin"""

from pathlib import Path
import os
import sys
import argparse
import rv32i_assembler
from rv32i_assembler import RV32IAssembler, AssemblerResults, Alert
from colorama import Fore, init # type: ignore

def get_program_files(source_dir):
    """Collects all the assembly program files from the source directory

    Args:
        source_dir (str): The name of hte source directory to search from
        
    Returns:
        list: Str list of all of the assembly programs in the dir
    """
    source_path = Path(source_dir)
    
    if not source_path.exists():
        print(f"Error: Source directory '{source_dir}' does not exist.")
        return []
    
    # Look for assembly files
    extensions = ['.s', '.asm', '.rv32i']
    programs = []
    
    for ext in extensions:
        programs.extend(source_path.glob(f"*{ext}"))

    return sorted([p.stem for p in programs])

def display_menu(programs):
    """Displays menu to terminal showing different programs available for assembly

    Args:
        programs (string list): List of valid program names
    """
    # ok, this repeating syntax might be my new favorite thing
    print("\n" + "="*50)
    print("RV32I Assembly Program Menu")
    print("="*50)
    
    if not programs:
        print('No assembly programs found in the source directory.')
        return
    
    for i, program in enumerate(programs, 1):
        print(f"{i:2d}. {program}")
        
    print("\nOptions:")
    print(f"Enter [1-{len(programs)}] to assemble a specific program")
    print("q - quit")
    print("="*50)

def assemble_program(assembler, source_dir, dest_dir, program_name):
    """Assembles a single program into binary form

    Args:
        assembler (RV32IAssembler): Assembler object to use
        source_dir (str): Source directory to search for the program
        dest_dir (str): Destination directory to place the binary file
        program_name (str): Name of the program to assemble

    Returns:
        dict: Dictionary describing the results of the assembly
    """
    source_path = Path(source_dir)
    dest_path = Path(dest_dir)
    
    # Ensure that destination actually exists
    dest_path.mkdir(parents=True, exist_ok=True)
    
    source_file = None
    extensions = ['.s', '.asm', '.rv32i']
    # Check for all extensions and name combinations
    for ext in extensions:
        potential_file = source_path / f"{program_name}{ext}"
        if potential_file.exists():
            source_file = potential_file
            break
        
        
    if not source_file:
        return {
            'program': program_name,
            'success': False,
            'error': f"Source file not found for '{program_name}'"
        }

    # Assembles to .mem files
    output_file = dest_path / f"{program_name}.mem"
    
    src =  open(source_file, 'r', encoding='ascii')
    dst =  open(output_file, 'w', encoding='ascii')
    
    
    # Warning flags handled in caller
    results = AssemblerResults()
    
    results = assembler.assemble(src, dst)
    
    src.close()
    dst.close() # in the future, the dst file shouldn't be created unless the compilation was successful
    
    return {
        'program': program_name,
        'success': results.success,
        'source': str(source_file),
        'output': str(output_file),
        'results': results
    }

def print_assembly_result(results):
    """Prints the results of the assembler given the return object

    Args:
        results (AssemblerResults): Dict containing the results of assembly of a program

    Raises:
        ValueError: Called if the assembler experiences an interna;/unknown error
    """
    
    print("\n" + "-"*50)
    print("ASSEMBLY RESULTS")
    print("-"*50)

    success = "SUCCESS" if results['success'] == True else "FAILURE"
    
    if results['success']:
        print(Fore.GREEN + f"\nProgram was successfully assembled with {len(results['results'].alerts)} warnings\n")
    else:
        print(Fore.RED + f"\nProgram failed assembly with {len(results['results'].alerts)} warnings and errors\n")
    
    # List all warnings and errors and stuff
    for idx, alert in enumerate(results["results"].alerts, 1):
        type = alert.type
        if type == 'internal':
            print("Internal/unknown fatal error occurred during assembly...")
            raise ValueError
        else:
            if type == "error":
                color = Fore.RED
            else:
                color = Fore.MAGENTA                
            
            print(color + f"{str(alert.type).capitalize()}", end=' ')
            print(f"on line {alert.line_num}.", end=' ')
            if type == 'warning':
                print(f"Warning type is: {alert.warning_type}.", end= ' ')
            print("\n" + str(alert.message), end="\n\n")
            
    print("-"*50)

def main():
    
    init(autoreset=True)
    
    parser = argparse.ArgumentParser(
        description="Semi-automation of RV32I program assembler",
        formatter_class=argparse.RawDescriptionHelpFormatter
    )
    
    parser.add_argument(
        '--source_dir', '-s',
        default='prog',
        help='Source directory containing assembly programs (default: prog)'
    )
    
    parser.add_argument(
        '--dest_dir', '-d',
        default='bin',
        help='Destination directory for assembled binaries (default: bin)'
    )
    
    parser.add_argument(
        '-A', '--assemble-all',
        action='store_true',
        help='Autoamtically assemble all programs and exit'
    )
    
    parser.add_argument(
        '-S', '--assemble-simple',
        action='store_true',
        help='Automatically assemble all programs with "simple_" prefix and exit'
    )
    
    # Add GCC style warnings
    parser.add_argument(
        '-Wall',
        action='store_true',
        help='Enable all warnings (like gcc -Wall)'
    )
    
    parser.add_argument(
        '-Werror',
        action='store_true',
        help='Treat all warnign as errors (like gcc -Werror)'
    )
    
    parser.add_argument(
        '-Wextra',
        action='store_true',
        help='Enable extra warnings (like gcc -Wextra)'
    )
    
    parser.add_argument(
        '-w',
        action='store_true',
        help='Suppress all warnings (like gcc -w)'
    )
    
    parser.add_argument(
        '-Wpedantic',
        action='store_true',
        help='Enable pedantic warnings for strict RISC-V compliance'
    )
    
    parser.add_argument(
        '-Wno-unused-label',
        action='store_true',
        help='Suppress warnings about unused labels'
    )
    
    parser.add_argument(
        '-Wno-immediate-range',
        action='store_true',
        help='Suppress warnings about immediate value ranges'
    )
    
    args = parser.parse_args()
    
    # Warning flag dictionary
    warning_flags = {
        'error_on_warning': args.Werror,
        'all_warnings': args.Wall,
        'extra_warnings': args.Wextra,
        'suppress_all': args.w,
        'pedantic': args.Wpedantic,
        'no_unused_label': args.Wno_unused_label,
        'no_immediate_range': args.Wno_immediate_range
    }
    
    # Check for conflicting flags
    if args.w and (args.Wall or args.Wextra or args.Wpedantic):
        print("Warning: -w flag suppresses all warnings, other warning flags will be ignored")
    
    try:
        assembler = RV32IAssembler(warning_Flags=warning_flags)
    except Exception as e:
        print(f"Error initializing assembler: {e}")
        sys.exit(1)
        
    programs = get_program_files(args.source_dir)
    
    if not programs:
        print(f"No assembly programs found in '{args.source_dir}' directory")
        sys.exit(1)
        
    if args.assemble_all:
        print(f"Assembling all {len(programs)} programs...")
        
        # todo: finish these
        sys.exit(0)
        
    if args.assemble_simple:
        simple_programs = [p for p in programs if p.startswith('simple_')]
        if not simple_programs:
            print("No programs with 'simple_' prefix found.")
            sys.exit(1)
            
        # todo: finish these
            
        sys.exit(0)

    print(f"Found {len(programs)} assembly programs in '{args.source_dir}'")
    
    while True:
        
        display_menu(programs)
        
        try:
            choice = input("\nEnter your choice: ").strip()

            if choice.lower() == 'q':
                print('\n\nGoodbye!')
                break
            
            try:
                choice_num = int(choice)            
                
                if 1 <= choice_num < len(programs):
                    program_name = programs[choice_num - 1]
                    
                    res = assemble_program(assembler, args.source_dir, args.dest_dir, program_name)
                    print_assembly_result(res)
                else:
                    print(f"Invalid choice. Please enter a number between 1 and {len(programs)}, or 'q' to quit.")
                
            except ValueError as ve:
                print(f"Invalid choice. Please enter a number between 1 and {len(programs)}, or 'q' to quit.")
                print(f"{ve}")
                
        except KeyboardInterrupt:
            print("\n\nGoodbye!")
            break
        except EOFError:
            print("\n\nGoodbye!")
            break


if __name__ == "__main__":
    main()