add support for using a btree database (DefinesDB) to store defines for preprocessing

The btree module, which ships with MicroPython, can efficiency manage a large
number of key-value pairs with minimal memory. It automatically initialises
to appropriate memory and cache limits, based on the device it's running
on, but if needed those parameters can be tuned too, e.g. to restrict memory
usage further.

The database is optional and must be supplied to the Preprocessor via the
use_db() method. It's safe however to always supply it, because a non-existing
database will behave like an empty database.

Care is taken not to unnecessarily create an empty db, when only reading from it
and not to unnecessarily check the file-system whether the database exists.

Inside the Preprocessor the database is opened and closed with a context manager.
This ensures the database will be closed properly again. While DefinesDB opens
the underlying database automatically, it cannot automatically close the database
again (using a destructor __del__ does not work, and MicroPython does not have
the "atexit" exit handler on the esp32). By using a context manager, the code
becomes cleaner, while still ensuring the database is closed at the end.

Author: wnienhaus

esp32_ulp/definesdb.py

@@ -0,0 +1,75 @@
+import os
+import btree
+from .util import file_exists
+
+DBNAME = 'defines.db'
+
+
+class DefinesDB:
+    def __init__(self):
+        self._file = None
+        self._db = None
+        self._db_exists = None
+
+    def clear(self):
+        self.close()
+        try:
+            os.remove(DBNAME)
+            self._db_exists = False
+        except OSError:
+            pass
+
+    def open(self):
+        if self._db:
+            return
+        try:
+            self._file = open(DBNAME, 'r+b')
+        except OSError:
+            self._file = open(DBNAME, 'w+b')
+        self._db = btree.open(self._file)
+        self._db_exists = True
+
+    def close(self):
+        if not self._db:
+            return
+        self._db.close()
+        self._db = None
+        self._file.close()
+        self._file = None
+
+    def db_exists(self):
+        if self._db_exists is None:
+            self._db_exists = file_exists(DBNAME)
+        return self._db_exists
+
+    def update(self, dictionary):
+        for k, v in dictionary.items():
+            self.__setitem__(k, v)
+
+    def get(self, key, default):
+        try:
+            result = self.__getitem__(key)
+        except KeyError:
+            result = default
+        return result
+
+    def keys(self):
+        if not self.db_exists():
+            return []
+
+        self.open()
+        return [k.decode() for k in self._db.keys()]
+
+    def __getitem__(self, key):
+        if not self.db_exists():
+            raise KeyError
+
+        self.open()
+        return self._db[key.encode()].decode()
+
+    def __setitem__(self, key, value):
+        self.open()
+        self._db[key.encode()] = str(value).encode()
+
+    def __iter__(self):
+        return iter(self.keys())

esp32_ulp/preprocess.py

@@ -1,5 +1,6 @@
 from . import nocomment
 from .util import split_tokens
+from .definesdb import DefinesDB
 
 
 class RTC_Macros:
@@ -56,6 +57,7 @@ def parse_define_line(self, line):
     def parse_defines(self, content):
         for line in content.splitlines():
             self._defines.update(self.parse_define_line(line))
+
         return self._defines
 
     def expand_defines(self, line):
@@ -66,21 +68,22 @@ def expand_defines(self, line):
             line = ""
             for t in tokens:
                 lu = self._defines.get(t, t)
+                if lu == t and self._defines_db:
+                    lu = self._defines_db.get(t, t)
                 if lu != t:
                     found = True
                 line += lu
 
         return line
 
     def process_include_file(self, filename):
-        defines = self._defines
-
-        with open(filename, 'r') as f:
-            for line in f:
-                result = self.parse_defines(line)
-                defines.update(result)
+        with self.open_db() as db:
+            with open(filename, 'r') as f:
+                for line in f:
+                    result = self.parse_define_line(line)
+                    db.update(result)
 
-        return defines
+        return db
 
     def expand_rtc_macros(self, line):
         clean_line = line.strip()
@@ -103,17 +106,43 @@ def expand_rtc_macros(self, line):
 
         return macro_fn(*macro_args)
 
+    def use_db(self, defines_db):
+        self._defines_db = defines_db
+
+    def open_db(self):
+        class ctx:
+            def __init__(self, db):
+                self._db = db
+
+            def __enter__(self):
+                # not opening DefinesDB - it opens itself when needed
+                return self._db
+
+            def __exit__(self, type, value, traceback):
+                if isinstance(self._db, DefinesDB):
+                    self._db.close()
+
+        if self._defines_db:
+            return ctx(self._defines_db)
+
+        return ctx(self._defines)
+
     def preprocess(self, content):
         self.parse_defines(content)
-        lines = nocomment.remove_comments(content)
-        result = []
-        for line in lines:
-            line = self.expand_defines(line)
-            line = self.expand_rtc_macros(line)
-            result.append(line)
-        result = "\n".join(result)
+
+        with self.open_db():
+            lines = nocomment.remove_comments(content)
+            result = []
+            for line in lines:
+                line = self.expand_defines(line)
+                line = self.expand_rtc_macros(line)
+                result.append(line)
+            result = "\n".join(result)
+
         return result
 
 
-def preprocess(content):
-    return Preprocessor().preprocess(content)
+def preprocess(content, use_defines_db=True):
+    preprocessor = Preprocessor()
+    preprocessor.use_db(DefinesDB())
+    return preprocessor.preprocess(content)

esp32_ulp/util.py

@@ -1,6 +1,7 @@
 DEBUG = False
 
 import gc
+import os
 
 NORMAL, WHITESPACE = 0, 1
 
@@ -67,3 +68,12 @@ def validate_expression(param):
                 if c not in '0123456789abcdef':
                     state = 0
     return True
+
+
+def file_exists(filename):
+    try:
+        os.stat(filename)
+        return True
+    except OSError:
+        pass
+    return False

tests/00_unit_tests.sh

@@ -4,7 +4,7 @@
 
 set -e
 
-for file in opcodes assemble link util preprocess; do
+for file in opcodes assemble link util preprocess definesdb; do
     echo testing $file...
     micropython $file.py
 done

tests/definesdb.py

@@ -0,0 +1,60 @@
+import os
+
+from esp32_ulp.definesdb import DefinesDB, DBNAME
+from esp32_ulp.util import file_exists
+
+tests = []
+
+
+def test(param):
+    tests.append(param)
+
+
+@test
+def test_definesdb_clear_removes_all_keys():
+    db = DefinesDB()
+    db.open()
+    db.update({'KEY1': 'VALUE1'})
+
+    db.clear()
+
+    assert 'KEY1' not in db
+
+    db.close()
+
+
+@test
+def test_definesdb_persists_data_across_instantiations():
+    db = DefinesDB()
+    db.open()
+    db.clear()
+
+    db.update({'KEY1': 'VALUE1'})
+
+    assert 'KEY1' in db
+
+    db.close()
+    del db
+    db = DefinesDB()
+    db.open()
+
+    assert db.get('KEY1', None) == 'VALUE1'
+
+    db.close()
+
+
+@test
+def test_definesdb_should_not_create_a_db_file_when_only_reading():
+    db = DefinesDB()
+
+    db.clear()
+    assert not file_exists(DBNAME)
+
+    assert db.get('some-key', None) is None
+    assert not file_exists(DBNAME)
+
+
+if __name__ == '__main__':
+    # run all methods marked with @test
+    for t in tests:
+        t()

tests/preprocess.py

@@ -1,4 +1,8 @@
+import os
+
 from esp32_ulp.preprocess import Preprocessor
+from esp32_ulp.definesdb import DefinesDB, DBNAME
+from esp32_ulp.util import file_exists
 
 tests = []
 
@@ -186,6 +190,7 @@ def test_process_include_file():
     p = Preprocessor()
 
     defines = p.process_include_file('fixtures/incl.h')
+
     assert defines['CONST1'] == '42'
     assert defines['CONST2'] == '99'
     assert defines.get('MULTI_LINE', None) == 'abc \\'  # correct. line continuations not supported
@@ -204,6 +209,94 @@ def test_process_include_file_with_multiple_files():
     assert defines['CONST3'] == '777', "constant from incl2.h"
 
 
+@test
+def test_process_include_file_using_database():
+    db = DefinesDB()
+    db.clear()
+
+    p = Preprocessor()
+    p.use_db(db)
+
+    p.process_include_file('fixtures/incl.h')
+    p.process_include_file('fixtures/incl2.h')
+
+    assert db['CONST1'] == '42', "constant from incl.h"
+    assert db['CONST2'] == '123', "constant overridden by incl2.h"
+    assert db['CONST3'] == '777', "constant from incl2.h"
+
+    db.close()
+
+
+@test
+def test_process_include_file_should_not_load_database_keys_into_instance_defines_dictionary():
+    db = DefinesDB()
+    db.clear()
+
+    p = Preprocessor()
+    p.use_db(db)
+
+    p.process_include_file('fixtures/incl.h')
+
+    # a bit hackish to reference instance-internal state
+    # but it's important to verify this, as we otherwise run out of memory on device
+    assert 'CONST2' not in p._defines
+
+
+
+@test
+def test_preprocess_should_use_definesdb_when_provided():
+    p = Preprocessor()
+
+    content = """\
+#define LOCALCONST 42
+
+entry:
+    move r1, LOCALCONST
+    move r2, DBKEY
+"""
+
+    # first try without db
+    result = p.preprocess(content)
+
+    assert "move r1, 42" in result
+    assert "move r2, DBKEY" in result
+    assert "move r2, 99" not in result
+
+    # now try with db
+    db = DefinesDB()
+    db.clear()
+    db.update({'DBKEY': '99'})
+    p.use_db(db)
+
+    result = p.preprocess(content)
+
+    assert "move r1, 42" in result
+    assert "move r2, 99" in result
+    assert "move r2, DBKEY" not in result
+
+
+@test
+def test_preprocess_should_ensure_no_definesdb_is_created_when_only_reading_from_it():
+    content = """\
+    #define CONST 42
+    move r1, CONST"""
+
+    # remove any existing db
+    db = DefinesDB()
+    db.clear()
+    assert not file_exists(DBNAME)
+
+    # now preprocess using db
+    p = Preprocessor()
+    p.use_db(db)
+
+    result = p.preprocess(content)
+
+    assert "move r1, 42" in result
+
+    assert not file_exists(DBNAME)
+
+
 if __name__ == '__main__':
     # run all methods marked with @test
     for t in tests: