Save programs on Prime Hub as they get downloaded

CHANGELOG.md

@@ -5,6 +5,7 @@
 ## [Unreleased]
 
 ### Added
+- Enabled `ujson` module.
 - Added ability to use more than one `DriveBase` in the same script.
 - Added `Motor.busy()` and `Motor.stalled()` methods, which
   (in case of `Motor`) are shorthand for `not Motor.control.done()` and

bricks/stm32/configport.h

@@ -61,6 +61,7 @@
 #define MICROPY_PY_ALL_SPECIAL_METHODS      (PYBRICKS_STM32_OPT_EXTRA_MOD)
 #define MICROPY_PY_REVERSE_SPECIAL_METHODS  (PYBRICKS_STM32_OPT_EXTRA_MOD)
 #define MICROPY_PY_IO                       (PYBRICKS_STM32_OPT_EXTRA_MOD)
+#define MICROPY_PY_UJSON                    (PYBRICKS_STM32_OPT_EXTRA_MOD)
 #define MICROPY_PY_STRUCT                   (PYBRICKS_STM32_OPT_EXTRA_MOD)
 #define MICROPY_PY_SYS                      (PYBRICKS_STM32_OPT_EXTRA_MOD)
 #define MICROPY_PY_SYS_EXIT         (0)

bricks/stm32/install_pybricks.py

@@ -40,17 +40,32 @@
 ERROR_BAD_FIRMWARE = """Invalid firmware file."""
 ERROR_BAD_HUB = """Incorrect hub type."""
 ERROR_DUAL_BOOT = """You are running dual-boot. Please re-install the official firmware first."""
+ERROR_INCOMPATIBLE_FIRMWARE = """The detected firmware is not compatible with this installer."""
 ERROR_EXTERNAL_FLASH = """Unable to create space for Pybricks firmware."""
 ERROR_FIRMWARE_COPY_FAILED = """Unable to copy the Pybricks firmware."""
 
-# Constants.
+# Flash constants.
 FLASH_FIRMWARE_START = 0x8008000
 FLASH_READ_SIZE = 32
 FLASH_WRITE_SIZE = FLASH_READ_SIZE * 4
+
+# Hub IDs.
+HUB_ID_SPIKE_PRIME = 0x81
+HUB_ID_SPIKE_ESSENTIAL = 0x83
 HUB_IDS = {
-    "LEGO Technic Large Hub(0x0009)": 0x81,
-    "LEGO Technic Large Hub(0x0010)": 0x81,
-    "LEGO Technic Small Hub(0x000D)": 0x83,
+    "LEGO Technic Large Hub(0x0009)": HUB_ID_SPIKE_PRIME,
+    "LEGO Technic Large Hub(0x0010)": HUB_ID_SPIKE_PRIME,
+    "LEGO Technic Small Hub(0x000D)": HUB_ID_SPIKE_ESSENTIAL,
+}
+
+# LEGO firmware versions on which Pybricks installation was successfully tested.
+KNOWN_SPIKE_PRIME_FW_VERSIONS = {
+    "v1.3.00.0000-e8c274a": "SPIKE App v2.0.4",
+    "v1.4.01.0000-594ce3d": "MINDSTORMS Robot Inventor App v10.3.0",
+}
+KNOWN_SPIKE_ESSENTIAL_FW_VERSIONS = {
+    "v1.0.00.0000-d94a557": "SPIKE App v2.0.0",
+    "v1.0.00.0070-51a2ff4": "SPIKE App v2.0.4",
 }
 
 
@@ -75,8 +90,8 @@ def get_lego_firmware_info():
     """Gets information about the running firmware."""
 
     # Get firmware/device ID
-    hub_id = firmware.id_string()
-    if hub_id not in HUB_IDS:
+    id_string = firmware.id_string()
+    if id_string not in HUB_IDS:
         stop_installation(ERROR_BAD_HUB)
 
     # Get firmware reset vector
@@ -90,7 +105,7 @@ def get_lego_firmware_info():
     firmware_version_position = read_internal_flash_int(FLASH_FIRMWARE_START + 0x200)
     firmware_version = read_internal_flash(firmware_version_position, 20).decode()
 
-    return hub_id, firmware_version, firmware_size, firmware_reset_vector
+    return id_string, firmware_version, firmware_size, firmware_reset_vector
 
 
 def get_lego_firmware(size):
@@ -134,7 +149,7 @@ def get_file_hash(path):
 def install(auto_reboot=True):
 
     # Get information about the running original LEGO firmware.
-    hub_id, lego_version, lego_size, lego_reset_vector = get_lego_firmware_info()
+    id_string, lego_version, lego_size, lego_reset_vector = get_lego_firmware_info()
 
     # Exit if user is running (outdated) Pybricks dual-boot firmware.
     if lego_reset_vector >= 0x80C0000:
@@ -145,6 +160,15 @@ def install(auto_reboot=True):
     print("    Version:", lego_version)
     print("    Size:", lego_size / 1024, "KB")
 
+    if (
+        HUB_IDS[id_string] == HUB_ID_SPIKE_PRIME
+        and lego_version not in KNOWN_SPIKE_PRIME_FW_VERSIONS
+    ) or (
+        HUB_IDS[id_string] == HUB_ID_SPIKE_ESSENTIAL
+        and lego_version not in KNOWN_SPIKE_ESSENTIAL_FW_VERSIONS
+    ):
+        stop_installation(ERROR_INCOMPATIBLE_FIRMWARE)
+
     # Back up the LEGO firmware so we can restore it from Pybricks later.
     print("Creating backup at:", LEGO_FIRMWARE_PATH)
     with open(LEGO_FIRMWARE_PATH, "wb") as backup_file:
@@ -156,7 +180,7 @@ def install(auto_reboot=True):
     with open(LEGO_METADATA_PATH, "w") as lego_meta_file:
         lego_backup_size, lego_backup_hash = get_file_hash(LEGO_FIRMWARE_PATH)
         lego_info = {
-            "device-id": HUB_IDS[hub_id],
+            "device-id": HUB_IDS[id_string],
             "firmware-sha256": lego_backup_hash,
             "firmware-size": lego_size,
             "firmware-version": lego_version,
@@ -175,7 +199,7 @@ def install(auto_reboot=True):
             stop_installation(ERROR_BAD_FIRMWARE)
 
         # Check that this firmware is made for this hub.
-        if HUB_IDS[hub_id] != pybricks_info["device-id"]:
+        if HUB_IDS[id_string] != pybricks_info["device-id"]:
             stop_installation(ERROR_BAD_HUB)
 
     # Display Pybricks firmware information.

bricks/stm32/main.c

@@ -14,6 +14,7 @@
 
 #include <pybricks/common.h>
 #include <pybricks/util_mp/pb_obj_helper.h>
+#include <pybricks/util_pb/pb_flash.h>
 
 #include "shared/readline/readline.h"
 #include "shared/runtime/gchelper.h"
@@ -185,8 +186,32 @@ static uint32_t get_user_program(uint8_t **buf, uint32_t *free_len) {
 
     // If button was pressed, return code to run script in flash
     if (err == PBIO_ERROR_CANCELED) {
+        #if (PYBRICKS_HUB_PRIMEHUB || PYBRICKS_HUB_ESSENTIALHUB)
+        // Open existing main.mpy file from flash
+        uint32_t size = 0;
+        if (pb_flash_file_open_get_size("/_pybricks/main.mpy", &size) != PBIO_SUCCESS) {
+            return 0;
+        }
+        // Check size and allocate buffer
+        if (size > MPY_MAX_BYTES) {
+            return 0;
+        }
+        *buf = m_malloc(size);
+        if (*buf == NULL) {
+            return 0;
+        }
+        // Read the file contents
+        if (pb_flash_file_read(*buf, size) != PBIO_SUCCESS) {
+            m_free(*buf);
+            return 0;
+        }
+        *free_len = size;
+        return size;
+        #else
+        // Load main program embedded in firmware
         *buf = &_pb_user_mpy_data;
         return _pb_user_mpy_size;
+        #endif
     }
 
     // Handle other errors
@@ -220,6 +245,12 @@ static uint32_t get_user_program(uint8_t **buf, uint32_t *free_len) {
     }
 
     *free_len = len;
+
+    #if (PYBRICKS_HUB_PRIMEHUB || PYBRICKS_HUB_ESSENTIALHUB)
+    // Save program as file
+    pb_flash_file_write("/_pybricks/main.mpy", *buf, len);
+    #endif // (PYBRICKS_HUB_PRIMEHUB || PYBRICKS_HUB_ESSENTIALHUB)
+
     return len;
 }
 
@@ -328,6 +359,12 @@ static void run_user_program(uint32_t len, uint8_t *buf, uint32_t free_len) {
 }
 
 static void stm32_main(void) {
+
+    #if (PYBRICKS_HUB_PRIMEHUB || PYBRICKS_HUB_ESSENTIALHUB)
+    mp_hal_delay_ms(500);
+    pb_flash_init();
+    #endif
+
 soft_reset:
     // Stack limit should be less than real stack size, so we have a chance
     // to recover from limit hit.  (Limit is measured in bytes.)

bricks/stm32/stm32.mk

@@ -264,6 +264,7 @@ PYBRICKS_PYBRICKS_SRC_C = $(addprefix pybricks/,\
 	util_pb/pb_conversions.c \
 	util_pb/pb_device_stm32.c \
 	util_pb/pb_error.c \
+	util_pb/pb_flash.c \
 	util_pb/pb_imu.c \
 	util_pb/pb_task.c \
 	)
@@ -565,7 +566,7 @@ ifeq ($(PB_LIB_BTSTACK),1)
 OBJ += $(addprefix $(BUILD)/, $(BTSTACK_SRC_C:.c=.o))
 endif
 ifeq ($(PB_LIB_LITTLEFS),1)
-CFLAGS+= -DLFS_NO_ASSERT -DLFS_NO_MALLOC -DLFS_NO_DEBUG -DLFS_NO_WARN -DLFS_NO_ERROR -DLFS_READONLY
+CFLAGS+= -DLFS_NO_ASSERT -DLFS_NO_MALLOC -DLFS_NO_DEBUG -DLFS_NO_WARN -DLFS_NO_ERROR
 OBJ += $(addprefix $(BUILD)/, $(LITTLEFS_SRC_C:.c=.o))
 endif
 ifeq ($(PB_USE_HAL),1)

lib/pbio/include/pbio/util.h

@@ -102,6 +102,22 @@ void pbio_set_uint32_le(uint8_t *buf, uint32_t value) {
     buf[3] = value >> 24;
 }
 
+#ifndef DOXYGEN
+static inline
+#endif
+/**
+ * Packs 32-bit big endian value into buffer.
+ *
+ * @param [in]  buf     The buffer.
+ * @param [in]  value   The value.
+ */
+void pbio_set_uint32_be(uint8_t *buf, uint32_t value) {
+    buf[0] = value >> 24;
+    buf[1] = value >> 16;
+    buf[2] = value >> 8;
+    buf[3] = value;
+}
+
 bool pbio_uuid128_reverse_compare(const uint8_t *uuid1, const uint8_t *uuid2);
 void pbio_uuid128_reverse_copy(uint8_t *dst, const uint8_t *src);
 

lib/pbio/platform/essential_hub/platform.c

@@ -434,6 +434,33 @@ void DMA2_Stream0_IRQHandler(void) {
     pbdrv_adc_stm32_hal_handle_irq();
 }
 
+void HAL_SPI_MspInit(SPI_HandleTypeDef *hspi) {
+    if (hspi->Instance == SPI2) {
+        // External flash
+        GPIO_InitTypeDef gpio_init;
+
+        // /CS, active low
+        gpio_init.Pin = GPIO_PIN_12;
+        gpio_init.Mode = GPIO_MODE_OUTPUT_PP;
+        gpio_init.Pull = GPIO_NOPULL;
+        gpio_init.Speed = GPIO_SPEED_FREQ_VERY_HIGH;
+        HAL_GPIO_Init(GPIOB, &gpio_init);
+        HAL_GPIO_WritePin(GPIOB, GPIO_PIN_12, GPIO_PIN_SET);
+
+        // SPI2_SCK
+        gpio_init.Pin = GPIO_PIN_10;
+        gpio_init.Mode = GPIO_MODE_AF_PP;
+        gpio_init.Pull = GPIO_NOPULL;
+        gpio_init.Speed = GPIO_SPEED_FREQ_VERY_HIGH;
+        gpio_init.Alternate = GPIO_AF5_SPI2;
+        HAL_GPIO_Init(GPIOB, &gpio_init);
+
+        // SPI2_MISO | SPI2_MOSI
+        gpio_init.Pin = GPIO_PIN_2 | GPIO_PIN_3;
+        HAL_GPIO_Init(GPIOC, &gpio_init);
+    }
+}
+
 // USB
 
 void HAL_PCD_MspInit(PCD_HandleTypeDef *hpcd) {
@@ -567,7 +594,7 @@ void SystemInit(void) {
         RCC_AHB1ENR_GPIODEN | RCC_AHB1ENR_DMA1EN | RCC_AHB1ENR_DMA2EN;
     RCC->APB1ENR |= RCC_APB1ENR_USART2EN | RCC_APB1ENR_USART3EN | RCC_APB1ENR_UART5EN |
         RCC_APB1ENR_TIM2EN | RCC_APB1ENR_TIM3EN | RCC_APB1ENR_TIM4EN |
-        RCC_APB1ENR_I2C3EN | RCC_APB1ENR_FMPI2C1EN;
+        RCC_APB1ENR_I2C3EN | RCC_APB1ENR_FMPI2C1EN | RCC_APB1ENR_SPI2EN;
     RCC->APB2ENR |= RCC_APB2ENR_TIM8EN | RCC_APB2ENR_ADC1EN | RCC_APB2ENR_SYSCFGEN;
     RCC->AHB2ENR |= RCC_AHB2ENR_OTGFSEN;
 

lib/pbio/platform/prime_hub/platform.c

@@ -671,6 +671,30 @@ void HAL_SPI_MspInit(SPI_HandleTypeDef *hspi) {
         gpio_init.Alternate = GPIO_AF5_SPI1;
         HAL_GPIO_Init(GPIOA, &gpio_init);
     }
+    if (hspi->Instance == SPI2) {
+        // External flash
+        GPIO_InitTypeDef gpio_init;
+
+        // /CS, active low
+        gpio_init.Pin = GPIO_PIN_12;
+        gpio_init.Mode = GPIO_MODE_OUTPUT_PP;
+        gpio_init.Pull = GPIO_NOPULL;
+        gpio_init.Speed = GPIO_SPEED_FREQ_VERY_HIGH;
+        HAL_GPIO_Init(GPIOB, &gpio_init);
+        HAL_GPIO_WritePin(GPIOB, GPIO_PIN_12, GPIO_PIN_SET);
+
+        // SPI2_SCK
+        gpio_init.Pin = GPIO_PIN_13;
+        gpio_init.Mode = GPIO_MODE_AF_PP;
+        gpio_init.Pull = GPIO_NOPULL;
+        gpio_init.Speed = GPIO_SPEED_FREQ_VERY_HIGH;
+        gpio_init.Alternate = GPIO_AF5_SPI2;
+        HAL_GPIO_Init(GPIOB, &gpio_init);
+
+        // SPI2_MISO | SPI2_MOSI
+        gpio_init.Pin = GPIO_PIN_2 | GPIO_PIN_3;
+        HAL_GPIO_Init(GPIOC, &gpio_init);
+    }
 }
 
 void SPI1_IRQHandler(void) {
@@ -808,7 +832,7 @@ void SystemInit(void) {
     RCC->APB1ENR |= RCC_APB1ENR_USART2EN | RCC_APB1ENR_UART4EN | RCC_APB1ENR_UART5EN |
         RCC_APB1ENR_UART7EN | RCC_APB1ENR_UART8EN | RCC_APB1ENR_TIM2EN | RCC_APB1ENR_TIM3EN |
         RCC_APB1ENR_TIM4EN | RCC_APB1ENR_TIM5EN | RCC_APB1ENR_TIM6EN | RCC_APB1ENR_TIM12EN |
-        RCC_APB1ENR_I2C2EN | RCC_APB1ENR_DACEN;
+        RCC_APB1ENR_I2C2EN | RCC_APB1ENR_DACEN | RCC_APB1ENR_SPI2EN;
     RCC->APB2ENR |= RCC_APB2ENR_TIM1EN | RCC_APB2ENR_TIM8EN | RCC_APB2ENR_UART9EN |
         RCC_APB2ENR_UART10EN | RCC_APB2ENR_ADC1EN | RCC_APB2ENR_SPI1EN | RCC_APB2ENR_SYSCFGEN;
     RCC->AHB2ENR |= RCC_AHB2ENR_OTGFSEN;