In-app flash write, UF2 parser, remove picotool dependency, reboot command

cmd/picomap/main.go

@@ -13,7 +13,7 @@ import (
 	"time"
 
 	"github.com/theater/picomap/lib/client"
-	"github.com/theater/picomap/lib/picotool"
+	"github.com/theater/picomap/lib/uf2"
 )
 
 func main() {
@@ -197,9 +197,12 @@ func cmdLog(_ []string) error {
 }
 
 func cmdLoad(args []string) error {
+	fs := flag.NewFlagSet("load", flag.ExitOnError)
+	dryRun := fs.Bool("dry-run", false, "parse UF2 and log operations without flashing")
+	fs.Parse(args)
 	target := "all"
-	if len(args) > 0 {
-		target = args[0]
+	if fs.NArg() > 0 {
+		target = fs.Arg(0)
 	}
 
 	wd, err := os.Getwd()
@@ -212,11 +215,6 @@ func cmdLoad(args []string) error {
 		return err
 	}
 
-	devs, err := client.ListSerial()
-	if err != nil {
-		return err
-	}
-
 	allTargets := []struct {
 		name string
 		uf2  string
@@ -240,6 +238,10 @@ func cmdLoad(args []string) error {
 		return fmt.Errorf("unknown target %q", target)
 	}
 
+	devs, err := client.ListSerial()
+	if err != nil {
+		return err
+	}
 	if len(devs) < len(targets) {
 		return fmt.Errorf("need %d device(s), found %d", len(targets), len(devs))
 	}
@@ -281,47 +283,17 @@ func cmdLoad(args []string) error {
 	for i := range devices {
 		log := slog.With("serial", devices[i].serial)
 		wg.Go(func() {
-			log.Info("sending PICOBOOT")
-			c, err := client.NewSerial(devices[i].dev, 500*time.Millisecond)
-			if err != nil {
-				errs[i] = err
-				return
-			}
-			err = c.PICOBOOT()
-			c.Close()
-			if err != nil {
-				errs[i] = fmt.Errorf("PICOBOOT %s: %w", devices[i].serial, err)
-				return
-			}
-			log.Info("PICOBOOT sent")
-		})
-	}
-	wg.Wait()
-	for i, err := range errs {
-		if err != nil {
-			return fmt.Errorf("[%s] %w", devices[i].serial, err)
-		}
-	}
-
-	uf2s := make([]string, len(targets))
-	for i := range targets {
-		uf2s[i] = targets[i].uf2
-	}
-
-	for i := range devices {
-		log := slog.With("serial", devices[i].serial)
-		wg.Go(func() {
-			log.Info("loading", "uf2", devices[i].name)
-			errs[i] = picotool.Load(devices[i].uf2, devices[i].serial, 10*time.Second)
+			log.Info("flashing", "uf2", devices[i].name)
+			errs[i] = flashDevice(devices[i].dev, devices[i].uf2, *dryRun, log)
 			if errs[i] == nil {
-				log.Info("loaded", "uf2", devices[i].name)
+				log.Info("flashed", "uf2", devices[i].name)
 			}
 		})
 	}
 	wg.Wait()
 	for i, err := range errs {
 		if err != nil {
-			return fmt.Errorf("[%s] load: %w", devices[i].serial, err)
+			return fmt.Errorf("[%s] flash: %w", devices[i].serial, err)
 		}
 	}
 
@@ -329,6 +301,54 @@ func cmdLoad(args []string) error {
 	return nil
 }
 
+func flashDevice(dev, uf2Path string, dryRun bool, log *slog.Logger) error {
+	blocks, err := uf2.Parse(uf2Path)
+	if err != nil {
+		return fmt.Errorf("parse uf2: %w", err)
+	}
+
+	log.Info("parsed uf2", "blocks", len(blocks))
+
+	const sectorSize = 4096
+
+	if dryRun {
+		erased := make(map[uint32]bool)
+		for _, b := range blocks {
+			sector := b.Addr &^ (sectorSize - 1)
+			if !erased[sector] {
+				log.Info("erasing", "addr", fmt.Sprintf("%08x", sector))
+				erased[sector] = true
+			}
+			log.Info("writing", "addr", fmt.Sprintf("%08x", b.Addr), "len", len(b.Data))
+		}
+		return nil
+	}
+
+	c, err := client.NewSerial(dev, 5*time.Second)
+	if err != nil {
+		return err
+	}
+	defer c.Close()
+
+	erased := make(map[uint32]bool)
+	for _, b := range blocks {
+		sector := b.Addr &^ (sectorSize - 1)
+		if !erased[sector] {
+			if err := c.FlashErase(sector, sectorSize); err != nil {
+				return fmt.Errorf("erase %08x: %w", sector, err)
+			}
+			erased[sector] = true
+		}
+		if err := c.FlashWrite(b.Addr, b.Data); err != nil {
+			return fmt.Errorf("write %08x: %w", b.Addr, err)
+		}
+	}
+
+	log.Info("rebooting")
+	c.Reboot()
+	return nil
+}
+
 func findTestDevice() (string, error) {
 	devs, err := client.ListSerial()
 	if err != nil {