Fix the try-catch lowering hack

This patch moves the success tag tracking of the try-call ast to its annotation instead of CompilerContext. This fixes the lowering of nested try-catch statements. This fixes CPR-1705.
matter-labs · Jun 5, 2024 · b116c85 · b116c85
1 parent 48b9a94
commit b116c85
Show file tree

Hide file tree

Showing 5 changed files with 26 additions and 13 deletions.
diff --git a/libsolidity/ast/ASTAnnotations.h b/libsolidity/ast/ASTAnnotations.h
@@ -317,6 +317,13 @@ struct FunctionCallAnnotation: ExpressionAnnotation
 	FunctionCallKind kind = FunctionCallKind::Unset;
 	/// If true, this is the external call of a try statement.
 	bool tryCall = false;
+
+	// HACK!
+	// We track the success tag here for the `TryStatement` lowering. This is to avoid the redundant status check and
+	// the conditional jump. Such patterns can confuse the zksolc translator.
+	//
+	// uint32_t since Assembly::new[Push]Tag() asserts that the tag is 32 bits.
+	std::optional<uint32_t> tryCallSuccessTag;
 };
 
 }
diff --git a/libsolidity/codegen/CompilerContext.h b/libsolidity/codegen/CompilerContext.h
@@ -350,11 +350,6 @@ class CompilerContext
 
 	RevertStrings revertStrings() const { return m_revertStrings; }
 
-	// HACK!
-	// We track the success tag here for the `TryStatement` lowering. This is to avoid the redundant status check and
-	// the conditional jump. Such patterns can confuse the zksolc translator.
-	evmasm::AssemblyItem currTryCallSuccessTag{evmasm::AssemblyItemType::UndefinedItem};
-
 private:
 	/// Updates source location set in the assembly.
 	void updateSourceLocation();

diff --git a/libsolidity/codegen/ContractCompiler.cpp b/libsolidity/codegen/ContractCompiler.cpp
@@ -936,7 +936,10 @@ bool ContractCompiler::visit(TryStatement const& _tryStatement)
 	StackHeightChecker checker(m_context);
 	CompilerContext::LocationSetter locationSetter(m_context, _tryStatement);
 
-	compileExpression(_tryStatement.externalCall());
+	auto* externalCall = dynamic_cast<FunctionCall const*>(&_tryStatement.externalCall());
+	solAssert(externalCall && externalCall->annotation().tryCall, "");
+	compileExpression(*externalCall);
+
 	int const returnSize = static_cast<int>(_tryStatement.externalCall().annotation().type->sizeOnStack());
 
 	// Stack: [ return values] <success flag>
@@ -949,8 +952,11 @@ bool ContractCompiler::visit(TryStatement const& _tryStatement)
 
 	evmasm::AssemblyItem endTag = m_context.appendJumpToNew();
 
-	solAssert(m_context.currTryCallSuccessTag.type() == AssemblyItemType::Tag, "");
-	m_context << m_context.currTryCallSuccessTag;
+	auto& successTag = externalCall->annotation().tryCallSuccessTag;
+	solAssert(successTag, "");
+	m_context << AssemblyItem(AssemblyItemType::Tag, *successTag);
+	successTag.reset();
+
 	m_context.adjustStackOffset(returnSize);
 	{
 		// Success case.

diff --git a/libsolidity/codegen/ExpressionCompiler.cpp b/libsolidity/codegen/ExpressionCompiler.cpp
@@ -705,7 +705,8 @@ bool ExpressionCompiler::visit(FunctionCall const& _functionCall)
 		case FunctionType::Kind::External:
 		case FunctionType::Kind::DelegateCall:
 			_functionCall.expression().accept(*this);
-			appendExternalFunctionCall(function, arguments, _functionCall.annotation().tryCall);
+			appendExternalFunctionCall(
+				function, arguments, _functionCall.annotation().tryCall, &_functionCall.annotation());
 			break;
 		case FunctionType::Kind::BareCallCode:
 			solAssert(false, "Callcode has been removed.");
@@ -765,7 +766,8 @@ bool ExpressionCompiler::visit(FunctionCall const& _functionCall)
 				// If this is a try call, return "<address> 1" in the success case and
 				// "0" in the error case.
 				AssemblyItem errorCase = m_context.appendConditionalJump();
-				m_context.currTryCallSuccessTag = m_context.appendJumpToNew();
+				_functionCall.annotation().tryCallSuccessTag
+					= m_context.appendJumpToNew().data().convert_to<uint32_t>();
 				m_context.adjustStackOffset(1);
 				m_context << errorCase;
 			}
@@ -2257,7 +2259,8 @@ void ExpressionCompiler::appendShiftOperatorCode(Token _operator, Type const& _v
 void ExpressionCompiler::appendExternalFunctionCall(
 	FunctionType const& _functionType,
 	vector<ASTPointer<Expression const>> const& _arguments,
-	bool _tryCall
+	bool _tryCall,
+	FunctionCallAnnotation* _annotation
 )
 {
 	solAssert(
@@ -2535,8 +2538,9 @@ void ExpressionCompiler::appendExternalFunctionCall(
 
 	if (_tryCall)
 	{
+		solAssert(_annotation, "");
 		// Success branch will reach this, failure branch will directly jump to endTag.
-		m_context.currTryCallSuccessTag = m_context.appendJumpToNew();
+		_annotation->tryCallSuccessTag = m_context.appendJumpToNew().data().convert_to<uint32_t>();
 		m_context.adjustStackOffset(1);
 		m_context << endTag;
 	}

diff --git a/libsolidity/codegen/ExpressionCompiler.h b/libsolidity/codegen/ExpressionCompiler.h
@@ -108,7 +108,8 @@ class ExpressionCompiler: private ASTConstVisitor
 	void appendExternalFunctionCall(
 		FunctionType const& _functionType,
 		std::vector<ASTPointer<Expression const>> const& _arguments,
-		bool _tryCall
+		bool _tryCall,
+		FunctionCallAnnotation* _annotation = nullptr
 	);
 	/// Appends code that evaluates a single expression and moves the result to memory. The memory offset is
 	/// expected to be on the stack and is updated by this call.