From bb32df471f162bc1f62ef065bf497754d3b20f75 Mon Sep 17 00:00:00 2001
From: Jack Koenig <koenig@sifive.com>
Date: Thu, 11 Apr 2024 16:49:24 -0700
Subject: [PATCH] Optimize TruthTable.split and .merge to do nothing in common
 case

Most TruthTables have a default of all dont care. We can optimize this
common case by doing actual work in TruthTable.split and
TruthTable.merge.
---
 .../util/experimental/decode/TruthTable.scala | 29 +++++++++++++------
 .../experimental/decode}/TruthTableSpec.scala | 23 +++++++++++++--
 2 files changed, 41 insertions(+), 11 deletions(-)
 rename src/test/scala/{chiselTests/util/experimental => chisel3/util/experimental/decode}/TruthTableSpec.scala (83%)

diff --git a/src/main/scala/chisel3/util/experimental/decode/TruthTable.scala b/src/main/scala/chisel3/util/experimental/decode/TruthTable.scala
index a0abfcb7d8f..8d637ed169c 100644
--- a/src/main/scala/chisel3/util/experimental/decode/TruthTable.scala
+++ b/src/main/scala/chisel3/util/experimental/decode/TruthTable.scala
@@ -162,7 +162,13 @@ object TruthTable {
     def index(bitPat: BitPat, bpType: Char): Seq[Int] =
       bitPat.rawString.zipWithIndex.filter(_._1 == bpType).map(_._2)
 
-    Seq('1', '0', '?').flatMap(t => tableFilter(index(table.default, t)))
+    // We need to split if the default has a mix of values (no need to split if all ones, all zeros, or all ?)
+    val needToSplit = !(table.default.allDontCares || table.default.allZeros || table.default.allOnes)
+    if (needToSplit) {
+      Seq('1', '0', '?').flatMap(t => tableFilter(index(table.default, t)))
+    } else {
+      Seq(table -> (0 until table.default.width))
+    }
   }
 
   /** consume tables, merge it into single table with different default bits.
@@ -185,13 +191,18 @@ object TruthTable {
       .sortBy(_._2)
       .map(_._1)
       .mkString}")
-    TruthTable(
-      tables
-        .flatMap(_._1.table.map(_._1))
-        .map { key =>
-          key -> bitPat(tables.flatMap { case (table, indexes) => reIndex(key, table, indexes) })
-        },
-      bitPat(tables.flatMap { case (table, indexes) => table.default.rawString.zip(indexes) })
-    )
+    val needToMerge = tables.size > 1
+    if (needToMerge) {
+      TruthTable(
+        tables
+          .flatMap(_._1.table.map(_._1))
+          .map { key =>
+            key -> bitPat(tables.flatMap { case (table, indexes) => reIndex(key, table, indexes) })
+          },
+        bitPat(tables.flatMap { case (table, indexes) => table.default.rawString.zip(indexes) })
+      )
+    } else {
+      tables.head._1
+    }
   }
 }
diff --git a/src/test/scala/chiselTests/util/experimental/TruthTableSpec.scala b/src/test/scala/chisel3/util/experimental/decode/TruthTableSpec.scala
similarity index 83%
rename from src/test/scala/chiselTests/util/experimental/TruthTableSpec.scala
rename to src/test/scala/chisel3/util/experimental/decode/TruthTableSpec.scala
index b239fa38291..c1cf81e7480 100644
--- a/src/test/scala/chiselTests/util/experimental/TruthTableSpec.scala
+++ b/src/test/scala/chisel3/util/experimental/decode/TruthTableSpec.scala
@@ -1,6 +1,6 @@
 // SPDX-License-Identifier: Apache-2.0
 
-package chiselTests.util.experimental
+package chisel3.util.experimental.decode
 
 import chisel3._
 import chisel3.util.BitPat
@@ -9,7 +9,7 @@ import org.scalatest.flatspec.AnyFlatSpec
 
 class TruthTableSpec extends AnyFlatSpec {
   val table = TruthTable(
-    Map(
+    Seq(
       // BitPat("b000") -> BitPat("b0"),
       BitPat("b001") -> BitPat("b?"),
       BitPat("b010") -> BitPat("b?"),
@@ -123,4 +123,23 @@ class TruthTableSpec extends AnyFlatSpec {
                                 |""".stripMargin)
     )
   }
+
+  behavior.of("TruthTable.split")
+
+  it should "not change the TruthTable when the default only uses a single value" in {
+    val result = TruthTable.split(table)
+    assert(result == Seq(table -> Seq(0)))
+    // As an optimization, it should return the same literal object
+    assert(result.head._1 eq table)
+  }
+
+  behavior.of("TruthTable.merge")
+
+  it should "not change the TruthTable when merging a single one" in {
+    val tables = Seq(table -> Seq(0))
+    val result = TruthTable.merge(tables)
+    assert(result == table)
+    // As an optimization, it should return the same literal object
+    assert(result eq table)
+  }
 }