2 files changed, 60 insertions, 17 deletions

diff --git a/libevmasm/Assembly.cpp b/libevmasm/Assembly.cpp
index 597fdae1..42b923df 100644
--- a/libevmasm/Assembly.cpp
+++ b/libevmasm/Assembly.cpp
@@ -350,38 +350,65 @@ void Assembly::injectStart(AssemblyItem const& _i)
 
 Assembly& Assembly::optimise(bool _enable, bool _isCreation, size_t _runs)
 {
-	optimiseInternal(_enable, _isCreation, _runs);
+	OptimiserSettings settings;
+	settings.isCreation = _isCreation;
+	settings.runPeephole = true;
+	if (_enable)
+	{
+		settings.runDeduplicate = true;
+		settings.runCSE = true;
+		settings.runConstantOptimiser = true;
+	}
+	settings.expectedExecutionsPerDeployment = _runs;
+	optimiseInternal(settings);
 	return *this;
 }
 
-map<u256, u256> Assembly::optimiseInternal(bool _enable, bool _isCreation, size_t _runs)
+
+Assembly& Assembly::optimise(OptimiserSettings _settings)
 {
+	optimiseInternal(_settings);
+	return *this;
+}
+
+map<u256, u256> Assembly::optimiseInternal(OptimiserSettings _settings)
+{
+	// Run optimisation for sub-assemblies.
 	for (size_t subId = 0; subId < m_subs.size(); ++subId)
 	{
-		map<u256, u256> subTagReplacements = m_subs[subId]->optimiseInternal(_enable, false, _runs);
+		OptimiserSettings settings = _settings;
+		// Disable creation mode for sub-assemblies.
+		settings.isCreation = false;
+		map<u256, u256> subTagReplacements = m_subs[subId]->optimiseInternal(settings);
+		// Apply the replacements (can be empty).
 		BlockDeduplicator::applyTagReplacement(m_items, subTagReplacements, subId);
 	}
 
 	map<u256, u256> tagReplacements;
+	// Iterate until no new optimisation possibilities are found.
 	for (unsigned count = 1; count > 0;)
 	{
 		count = 0;
 
-		PeepholeOptimiser peepOpt(m_items);
-		while (peepOpt.optimise())
-			count++;
-
-		if (!_enable)
-			continue;
+		if (_settings.runPeephole)
+		{
+			PeepholeOptimiser peepOpt(m_items);
+			while (peepOpt.optimise())
+				count++;
+		}
 
 		// This only modifies PushTags, we have to run again to actually remove code.
-		BlockDeduplicator dedup(m_items);
-		if (dedup.deduplicate())
+		if (_settings.runDeduplicate)
 		{
-			tagReplacements.insert(dedup.replacedTags().begin(), dedup.replacedTags().end());
-			count++;
+			BlockDeduplicator dedup(m_items);
+			if (dedup.deduplicate())
+			{
+				tagReplacements.insert(dedup.replacedTags().begin(), dedup.replacedTags().end());
+				count++;
+			}
 		}
 
+		if (_settings.runCSE)
 		{
 			// Control flow graph optimization has been here before but is disabled because it
 			// assumes we only jump to tags that are pushed. This is not the case anymore with
@@ -429,10 +456,10 @@ map<u256, u256> Assembly::optimiseInternal(bool _enable, bool _isCreation, size_
 		}
 	}
 
-	if (_enable)
+	if (_settings.runConstantOptimiser)
 		ConstantOptimisationMethod::optimiseConstants(
-			_isCreation,
-			_isCreation ? 1 : _runs,
+			_settings.isCreation,
+			_settings.isCreation ? 1 : _settings.expectedExecutionsPerDeployment,
 			*this,
 			m_items
 		);
diff --git a/libevmasm/Assembly.h b/libevmasm/Assembly.h
index 0d40abcf..451b4ea0 100644
--- a/libevmasm/Assembly.h
+++ b/libevmasm/Assembly.h
@@ -97,12 +97,28 @@ public:
 	LinkerObject const& assemble() const;
 	bytes const& data(h256 const& _i) const { return m_data.at(_i); }
 
+	struct OptimiserSettings
+	{
+		bool isCreation = false;
+		bool runPeephole = false;
+		bool runDeduplicate = false;
+		bool runCSE = false;
+		bool runConstantOptimiser = false;
+		/// This specifies an estimate on how often each opcode in this assembly will be executed,
+		/// i.e. use a small value to optimise for size and a large value to optimise for runtime gas usage.
+		size_t expectedExecutionsPerDeployment = 200;
+	};
+
+	/// Execute optimisation passes as defined by @a _settings and return the optimised assembly.
+	Assembly& optimise(OptimiserSettings _settings);
+
 	/// Modify (if @a _enable is set) and return the current assembly such that creation and
 	/// execution gas usage is optimised. @a _isCreation should be true for the top-level assembly.
 	/// @a _runs specifes an estimate on how often each opcode in this assembly will be executed,
 	/// i.e. use a small value to optimise for size and a large value to optimise for runtime.
 	/// If @a _enable is not set, will perform some simple peephole optimizations.
 	Assembly& optimise(bool _enable, bool _isCreation = true, size_t _runs = 200);
+
 	Json::Value stream(
 		std::ostream& _out,
 		std::string const& _prefix = "",
@@ -113,7 +129,7 @@ public:
 protected:
 	/// Does the same operations as @a optimise, but should only be applied to a sub and
 	/// returns the replaced tags.
-	std::map<u256, u256> optimiseInternal(bool _enable, bool _isCreation, size_t _runs);
+	std::map<u256, u256> optimiseInternal(OptimiserSettings _settings);
 
 	unsigned bytesRequired(unsigned subTagSize) const;