2 files changed, 98 insertions, 25 deletions

diff --git a/passes/cmds/setundef.cc b/passes/cmds/setundef.cc
index f6949c820..3eedc86b8 100644
--- a/passes/cmds/setundef.cc
+++ b/passes/cmds/setundef.cc
@@ -393,44 +393,112 @@ struct SetundefPass : public Pass {
 						ffbits.insert(bit);
 				}
 
-				for (auto wire : module->wires())
+				auto process_initwires = [&]()
 				{
-					if (!wire->attributes.count("\\init"))
-						continue;
+					dict<Wire*, int> wire_weights;
 
-					for (auto bit : sigmap(wire))
-						ffbits.erase(bit);
+					for (auto wire : initwires)
+					{
+						int weight = 0;
 
-					initwires.insert(wire);
-				}
+						for (auto bit : sigmap(wire))
+							weight += ffbits.count(bit) ? +1 : -1;
+
+						wire_weights[wire] = weight;
+					}
+
+					initwires.sort([&](Wire *a, Wire *b) { return wire_weights.at(a) > wire_weights.at(b); });
+
+					for (auto wire : initwires)
+					{
+						Const &initval = wire->attributes["\\init"];
+						initval.bits.resize(GetSize(wire), State::Sx);
+
+						for (int i = 0; i < GetSize(wire); i++) {
+							SigBit bit = sigmap(SigBit(wire, i));
+							if (initval[i] == State::Sx && ffbits.count(bit)) {
+								initval[i] = worker.next_bit();
+								ffbits.erase(bit);
+							}
+						}
+
+						if (initval.is_fully_undef())
+							wire->attributes.erase("\\init");
+					}
+
+					initwires.clear();
+				};
 
 				for (int wire_types = 0; wire_types < 2; wire_types++)
-					for (auto wire : module->wires())
+				{
+					// prioritize wires that already have an init attribute
+					if (!ffbits.empty())
 					{
-						if (wire->name[0] == (wire_types ? '\\' : '$'))
-					next_wire:
-							continue;
+						for (auto wire : module->wires())
+						{
+							if (wire->name[0] == (wire_types ? '\\' : '$'))
+								continue;
 
-						for (auto bit : sigmap(wire))
-							if (!ffbits.count(bit))
-								goto next_wire;
+							if (!wire->attributes.count("\\init"))
+								continue;
 
-						for (auto bit : sigmap(wire))
-							ffbits.erase(bit);
+							Const &initval = wire->attributes["\\init"];
+							initval.bits.resize(GetSize(wire), State::Sx);
+
+							if (initval.is_fully_undef()) {
+								wire->attributes.erase("\\init");
+								continue;
+							}
 
-						initwires.insert(wire);
+							for (int i = 0; i < GetSize(wire); i++)
+								if (initval[i] != State::Sx)
+									ffbits.erase(sigmap(SigBit(wire, i)));
+
+							initwires.insert(wire);
+						}
+
+						process_initwires();
 					}
 
-				for (auto wire : initwires)
-				{
-					Const &initval = wire->attributes["\\init"];
+					// next consider wires that completely contain bits to be initialized
+					if (!ffbits.empty())
+					{
+						for (auto wire : module->wires())
+						{
+							if (wire->name[0] == (wire_types ? '\\' : '$'))
+								continue;
 
-					for (int i = 0; i < GetSize(wire); i++)
-						if (GetSize(initval) <= i)
-							initval.bits.push_back(worker.next_bit());
-						else if (initval.bits[i] == State::Sx)
-							initval.bits[i] = worker.next_bit();
+							for (auto bit : sigmap(wire))
+								if (!ffbits.count(bit))
+									goto next_wire;
+
+							initwires.insert(wire);
+
+						next_wire:
+							continue;
+						}
+
+						process_initwires();
+					}
+
+					// finally use whatever wire we can find.
+					if (!ffbits.empty())
+					{
+						for (auto wire : module->wires())
+						{
+							if (wire->name[0] == (wire_types ? '\\' : '$'))
+								continue;
+
+							for (auto bit : sigmap(wire))
+								if (ffbits.count(bit))
+									initwires.insert(wire);
+						}
+
+						process_initwires();
+					}
 				}
+
+				log_assert(ffbits.empty());
 			}
 
 			module->rewrite_sigspecs(worker);
diff --git a/passes/cmds/tee.cc b/passes/cmds/tee.cc
index ee96ace86..1a44bdaec 100644
--- a/passes/cmds/tee.cc
+++ b/passes/cmds/tee.cc
@@ -52,7 +52,9 @@ struct TeePass : public Pass {
 	void execute(std::vector<std::string> args, RTLIL::Design *design) YS_OVERRIDE
 	{
 		std::vector<FILE*> backup_log_files, files_to_close;
+		std::vector<std::ostream*> backup_log_streams;
 		int backup_log_verbose_level = log_verbose_level;
+		backup_log_streams = log_streams;
 		backup_log_files = log_files;
 
 		size_t argidx;
@@ -60,6 +62,7 @@ struct TeePass : public Pass {
 		{
 			if (args[argidx] == "-q" && files_to_close.empty()) {
 				log_files.clear();
+				log_streams.clear();
 				continue;
 			}
 			if ((args[argidx] == "-o" || args[argidx] == "-a") && argidx+1 < args.size()) {
@@ -89,6 +92,7 @@ struct TeePass : public Pass {
 			for (auto cf : files_to_close)
 				fclose(cf);
 			log_files = backup_log_files;
+			log_streams = backup_log_streams;
 			throw;
 		}
 
@@ -97,6 +101,7 @@ struct TeePass : public Pass {
 
 		log_verbose_level = backup_log_verbose_level;
 		log_files = backup_log_files;
+		log_streams = backup_log_streams;
 	}
 } TeePass;