From f6e02497ae063b7939b0670153d22a7b17bf4408 Mon Sep 17 00:00:00 2001
From: Simon Glass <sjg@chromium.org>
Date: Sat, 20 Jul 2019 12:24:08 -0600
Subject: [PATCH] binman: Update state when replacing device-tree entries

Since the state module holds references to all the device trees used by
binman, it must be updated when the device trees are updated. Add support
for this.

Signed-off-by: Simon Glass <sjg@chromium.org>
---
 tools/binman/etype/blob_dtb.py |  9 +++++++++
 tools/binman/state.py          | 16 ++++++++++++++++
 tools/dtoc/fdt.py              |  8 ++++++++
 tools/dtoc/test_fdt.py         |  5 +++++
 4 files changed, 38 insertions(+)

diff --git a/tools/binman/etype/blob_dtb.py b/tools/binman/etype/blob_dtb.py
index a3b548eef2..5b559967d7 100644
--- a/tools/binman/etype/blob_dtb.py
+++ b/tools/binman/etype/blob_dtb.py
@@ -53,3 +53,12 @@ class Entry_blob_dtb(Entry_blob):
         """
         fname = self.GetDefaultFilename()
         return {self.GetFdtEtype(): [self, fname]}
+
+    def WriteData(self, data, decomp=True):
+        ok = Entry_blob.WriteData(self, data, decomp)
+
+        # Update the state module, since it has the authoritative record of the
+        # device trees used. If we don't do this, then state.GetFdtContents()
+        # will still return the old contents
+        state.UpdateFdtContents(self.GetFdtEtype(), data)
+        return ok
diff --git a/tools/binman/state.py b/tools/binman/state.py
index f22cc82d87..d704ed2c7c 100644
--- a/tools/binman/state.py
+++ b/tools/binman/state.py
@@ -108,6 +108,22 @@ def GetFdtContents(etype='u-boot-dtb'):
         data = tools.ReadFile(pathname)
     return pathname, data
 
+def UpdateFdtContents(etype, data):
+    """Update the contents of a particular device tree
+
+    The device tree is updated and written back to its file. This affects what
+    is returned from future called to GetFdtContents(), etc.
+
+    Args:
+        etype: Entry type (e.g. 'u-boot-dtb')
+        data: Data to replace the DTB with
+    """
+    dtb, fname, entry = output_fdt_info[etype]
+    dtb_fname = dtb.GetFilename()
+    tools.WriteFile(dtb_fname, data)
+    dtb = fdt.FdtScan(dtb_fname)
+    output_fdt_info[etype] = [dtb, fname, entry]
+
 def SetEntryArgs(args):
     """Set the value of the entry args
 
diff --git a/tools/dtoc/fdt.py b/tools/dtoc/fdt.py
index cd7673c7da..6770be79fb 100644
--- a/tools/dtoc/fdt.py
+++ b/tools/dtoc/fdt.py
@@ -695,6 +695,14 @@ class Fdt:
         node = Node(fdt, parent, offset, name, path)
         return node
 
+    def GetFilename(self):
+        """Get the filename of the device tree
+
+        Returns:
+            String filename
+        """
+        return self._fname
+
 def FdtScan(fname):
     """Returns a new Fdt object"""
     dtb = Fdt(fname)
diff --git a/tools/dtoc/test_fdt.py b/tools/dtoc/test_fdt.py
index 16a4430892..028c8cbaa8 100755
--- a/tools/dtoc/test_fdt.py
+++ b/tools/dtoc/test_fdt.py
@@ -449,6 +449,11 @@ class TestProp(unittest.TestCase):
         self.assertIn("node '/spl-test': Missing property 'one'",
                       str(e.exception))
 
+    def testGetFilename(self):
+        """Test the dtb filename can be provided"""
+        self.assertEqual(tools.GetOutputFilename('source.dtb'),
+                         self.dtb.GetFilename())
+
 
 class TestFdtUtil(unittest.TestCase):
     """Tests for the fdt_util module