From 99c48d36ab88ccee2f326ed3727139b5ba2d9c1b Mon Sep 17 00:00:00 2001
From: Anton Lydike <me@antonlydike.de>
Date: Thu, 26 Jan 2023 17:45:07 +0000
Subject: [PATCH] add an unlimited register mode

---
 riscemu/__main__.py              | 12 ++++++++++--
 riscemu/config.py                |  1 +
 riscemu/instructions/__init__.py |  2 +-
 riscemu/parser.py                |  3 +++
 riscemu/registers.py             | 22 ++++++++++++++++++++--
 riscemu/types/cpu.py             |  2 +-
 6 files changed, 36 insertions(+), 6 deletions(-)

diff --git a/riscemu/__main__.py b/riscemu/__main__.py
index 1d46d31..8c4f8a7 100644
--- a/riscemu/__main__.py
+++ b/riscemu/__main__.py
@@ -57,12 +57,19 @@ if __name__ == '__main__':
             setattr(namespace, self.dest, d)
 
 
-    parser = argparse.ArgumentParser(description='RISC-V Userspace parser and emulator', prog='riscemu')
+    parser = argparse.ArgumentParser(description='RISC-V Userspace parser and emulator', prog='riscemu',
+                                     formatter_class=argparse.RawTextHelpFormatter)
     parser.add_argument('files', metavar='file.asm', type=str, nargs='+',
                         help='The assembly files to load, the last one will be run')
 
     parser.add_argument('--options', '-o', action=OptionStringAction,
-                        keys=('disable_debug', 'no_syscall_symbols', 'fail_on_ex', 'add_accept_imm'))
+                        keys=('disable_debug', 'no_syscall_symbols', 'fail_on_ex', 'add_accept_imm', 'unlimited_regs'),
+                        help="""Toggle options. Available options are:
+disable_debug:        Disable ebreak instructions
+no_syscall_symbols:   Don't add symbols for SCALL_EXIT and others
+fail_on_ex:           If set, exceptions won't trigger the debugger
+add_accept_imm:       Accept "add rd, rs, imm" instruction (instead of addi)
+unlimited_regs:       Allow an unlimited number of registers""")
 
     parser.add_argument('--syscall-opts', '-so', action=OptionStringAction,
                         keys=('fs_access', 'disable_input'))
@@ -88,6 +95,7 @@ if __name__ == '__main__':
         include_scall_symbols=not args.options['no_syscall_symbols'],
         debug_on_exception=not args.options['fail_on_ex'],
         add_accept_imm=args.options['add_accept_imm'],
+        unlimited_registers=args.options['unlimited_regs'],
         scall_fs=args.syscall_opts['fs_access'],
         scall_input=not args.syscall_opts['disable_input'],
         verbosity=args.verbose
diff --git a/riscemu/config.py b/riscemu/config.py
index a3dc5ae..9bda381 100644
--- a/riscemu/config.py
+++ b/riscemu/config.py
@@ -20,6 +20,7 @@ class RunConfig:
     scall_fs: bool = False
     verbosity: int = 0
     slowdown: float = 1
+    unlimited_registers: bool = False
 
 
 CONFIG = RunConfig()
diff --git a/riscemu/instructions/__init__.py b/riscemu/instructions/__init__.py
index 96fb524..3343ae8 100644
--- a/riscemu/instructions/__init__.py
+++ b/riscemu/instructions/__init__.py
@@ -6,7 +6,7 @@ SPDX-License-Identifier: MIT
 This package holds all instruction sets, available to the processor
 """
 
-from .instruction_set import InstructionSet
+from .instruction_set import InstructionSet, Instruction
 from .RV32M import RV32M
 from .RV32I import RV32I
 from .RV32A import RV32A
diff --git a/riscemu/parser.py b/riscemu/parser.py
index 83f960e..86759ec 100644
--- a/riscemu/parser.py
+++ b/riscemu/parser.py
@@ -110,6 +110,9 @@ class AssemblyFileLoader(ProgramLoader):
         with open(self.source_path, 'r') as f:
             return parse_tokens(self.filename, tokenize(f))
 
+    def parse_io(self, io):
+        return parse_tokens(self.filename, tokenize(io))
+
     @classmethod
     def can_parse(cls, source_path: str) -> float:
         """
diff --git a/riscemu/registers.py b/riscemu/registers.py
index cce48c4..80ebf6c 100644
--- a/riscemu/registers.py
+++ b/riscemu/registers.py
@@ -17,11 +17,22 @@ class Registers:
     Represents a bunch of registers
     """
 
-    def __init__(self):
+    valid_regs = {
+        'zero', 'ra', 'sp', 'gp', 'tp', 's0', 'fp',
+        't0', 't1', 't2', 't3', 't4', 't5', 't6',
+        's1', 's2', 's3', 's4', 's5', 's6', 's7', 's8', 's9', 's10', 's11',
+        'a0', 'a1', 'a2', 'a3', 'a4', 'a5', 'a6', 'a7',
+        'ft0', 'ft1', 'ft2', 'ft3', 'ft4', 'ft5', 'ft6', 'ft7',
+        'fs0', 'fs1', 'fs2', 'fs3', 'fs4', 'fs5', 'fs6', 'fs7', 'fs8', 'fs9', 'fs10', 'fs11',
+        'fa0', 'fa1', 'fa2', 'fa3', 'fa4', 'fa5', 'fa6', 'fa7'
+    }
+
+    def __init__(self, infinite_regs: bool = False):
         from .types import Int32
         self.vals = defaultdict(lambda: Int32(0))
         self.last_set = None
         self.last_read = None
+        self.infinite_regs = infinite_regs
 
     def dump(self, full=False):
         """
@@ -108,7 +119,10 @@ class Registers:
             reg = 's1'
         if mark_set:
             self.last_set = reg
-        # check 32 bit signed bounds
+
+        if not self.infinite_regs and reg not in self.valid_regs:
+            raise RuntimeError("Invalid register: {}".format(reg))
+
         self.vals[reg] = val.unsigned()
         return True
 
@@ -123,6 +137,10 @@ class Registers:
         #    raise InvalidRegisterException(reg)
         if reg == 'fp':
             reg = 's0'
+
+        if not self.infinite_regs and reg not in self.valid_regs:
+            raise RuntimeError("Invalid register: {}".format(reg))
+
         if mark_read:
             self.last_read = reg
         return self.vals[reg]
diff --git a/riscemu/types/cpu.py b/riscemu/types/cpu.py
index 75df6c3..4672f0c 100644
--- a/riscemu/types/cpu.py
+++ b/riscemu/types/cpu.py
@@ -31,7 +31,7 @@ class CPU(ABC):
 
     def __init__(self, mmu: 'MMU', instruction_sets: List[Type['InstructionSet']], conf: RunConfig):
         self.mmu = mmu
-        self.regs = Registers()
+        self.regs = Registers(conf.unlimited_registers)
         self.conf = conf
 
         self.instruction_sets = set()