/src/solidity/libyul/optimiser/StackToMemoryMover.cpp

Source
/*
  This file is part of solidity.

  solidity is free software: you can redistribute it and/or modify
  it under the terms of the GNU General Public License as published by
  the Free Software Foundation, either version 3 of the License, or
  (at your option) any later version.

  solidity is distributed in the hope that it will be useful,
  but WITHOUT ANY WARRANTY; without even the implied warranty of
  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  GNU General Public License for more details.

  You should have received a copy of the GNU General Public License
  along with solidity.  If not, see <http://www.gnu.org/licenses/>.
*/
#include <libyul/optimiser/StackToMemoryMover.h>
#include <libyul/optimiser/NameCollector.h>
#include <libyul/optimiser/NameDispenser.h>
#include <libyul/backends/evm/EVMDialect.h>

#include <libyul/AST.h>
#include <libyul/Utilities.h>

#include <libsolutil/CommonData.h>

#include <range/v3/algorithm/none_of.hpp>
#include <range/v3/view/filter.hpp>
#include <range/v3/view/transform.hpp>
#include <range/v3/view/zip.hpp>
#include <range/v3/range/conversion.hpp>

using namespace solidity;
using namespace solidity::yul;

namespace
{
std::vector<Statement> generateMemoryStore(
  Dialect const& _dialect,
  langutil::DebugData::ConstPtr const& _debugData,
  LiteralValue const& _mpos,
  Expression _value
)
{
  std::optional<BuiltinHandle> memoryStoreFunctionHandle = _dialect.memoryStoreFunctionHandle();
  yulAssert(memoryStoreFunctionHandle);
  std::vector<Statement> result;
  result.emplace_back(ExpressionStatement{_debugData, FunctionCall{
    _debugData,
    BuiltinName{_debugData, *memoryStoreFunctionHandle},
    {
      Literal{_debugData, LiteralKind::Number, _mpos},
      std::move(_value)
    }
  }});
  return result;
}

FunctionCall generateMemoryLoad(Dialect const& _dialect, langutil::DebugData::ConstPtr const& _debugData, LiteralValue const& _mpos)
{
  std::optional<BuiltinHandle> const& memoryLoadHandle = _dialect.memoryLoadFunctionHandle();
  yulAssert(memoryLoadHandle);
  return FunctionCall{
    _debugData,
    BuiltinName{_debugData, *memoryLoadHandle}, {
      Literal{
        _debugData,
        LiteralKind::Number,
        _mpos
      }
    }
  };
}
}

void StackToMemoryMover::run(
  OptimiserStepContext& _context,
  u256 _reservedMemory,
  std::map<YulName, uint64_t> const& _memorySlots,
  uint64_t _numRequiredSlots,
  Block& _block
)
{
  VariableMemoryOffsetTracker memoryOffsetTracker(_reservedMemory, _memorySlots, _numRequiredSlots);
  StackToMemoryMover stackToMemoryMover(
    _context,
    memoryOffsetTracker,
    util::applyMap(
      allFunctionDefinitions(_block),
      util::mapTuple([](YulName _name, FunctionDefinition const* _funDef) {
        return make_pair(_name, _funDef->returnVariables);
      }),
      std::map<YulName, NameWithDebugDataList>{}
    )
  );
  stackToMemoryMover(_block);
  _block.statements += std::move(stackToMemoryMover.m_newFunctionDefinitions);
}

StackToMemoryMover::StackToMemoryMover(
  OptimiserStepContext& _context,
  VariableMemoryOffsetTracker const& _memoryOffsetTracker,
  std::map<YulName, NameWithDebugDataList> _functionReturnVariables
):
m_context(_context),
m_memoryOffsetTracker(_memoryOffsetTracker),
m_nameDispenser(_context.dispenser),
m_functionReturnVariables(std::move(_functionReturnVariables))
{
  auto const* evmDialect = dynamic_cast<EVMDialect const*>(&_context.dialect);
  yulAssert(
    evmDialect && evmDialect->providesObjectAccess(),
    "StackToMemoryMover can only be run on objects using the EVMDialect with object access."
  );
}

void StackToMemoryMover::operator()(FunctionDefinition& _functionDefinition)
{
  // It is important to first visit the function body, so that it doesn't replace the memory inits for
  // variable arguments we might generate below.
  ASTModifier::operator()(_functionDefinition);

  std::vector<Statement> memoryVariableInits;

  // All function parameters with a memory slot are moved at the beginning of the function body.
  for (NameWithDebugData const& param: _functionDefinition.parameters)
    if (auto slot = m_memoryOffsetTracker(param.name))
      memoryVariableInits += generateMemoryStore(
        m_context.dialect,
        param.debugData,
        *slot,
        Identifier{param.debugData, param.name}
      );

  // All memory return variables have to be initialized to zero in memory.
  for (NameWithDebugData const& returnVariable: _functionDefinition.returnVariables)
    if (auto slot = m_memoryOffsetTracker(returnVariable.name))
      memoryVariableInits += generateMemoryStore(
        m_context.dialect,
        returnVariable.debugData,
        *slot,
        Literal{returnVariable.debugData, LiteralKind::Number, LiteralValue(u256{0})}
      );

  // Special case of a function with a single return argument that needs to move to memory.
  if (_functionDefinition.returnVariables.size() == 1 && m_memoryOffsetTracker(_functionDefinition.returnVariables.front().name))
  {
    NameWithDebugDataList stackParameters = _functionDefinition.parameters | ranges::views::filter(
      std::not_fn(m_memoryOffsetTracker)
    ) | ranges::to<NameWithDebugDataList>;
    // Generate new function without return variable and with only the non-moved parameters.
    YulName newFunctionName = m_context.dispenser.newName(_functionDefinition.name);
    m_newFunctionDefinitions.emplace_back(FunctionDefinition{
      _functionDefinition.debugData,
      newFunctionName,
      stackParameters,
      {},
      std::move(_functionDefinition.body)
    });
    // Generate new names for the arguments to maintain disambiguation.
    std::map<YulName, YulName> newArgumentNames;
    for (NameWithDebugData const& _var: stackParameters)
      newArgumentNames[_var.name] = m_context.dispenser.newName(_var.name);
    for (auto& parameter: _functionDefinition.parameters)
      parameter.name = util::valueOrDefault(newArgumentNames, parameter.name, parameter.name);
    // Replace original function by a call to the new function and an assignment to the return variable from memory.
    _functionDefinition.body = Block{_functionDefinition.debugData, std::move(memoryVariableInits)};
    _functionDefinition.body.statements.emplace_back(ExpressionStatement{
      _functionDefinition.debugData,
      FunctionCall{
        _functionDefinition.debugData,
        Identifier{_functionDefinition.debugData, newFunctionName},
        stackParameters | ranges::views::transform([&](NameWithDebugData const& _arg) {
          return Expression{Identifier{_arg.debugData, newArgumentNames.at(_arg.name)}};
        }) | ranges::to<std::vector<Expression>>
      }
    });
    _functionDefinition.body.statements.emplace_back(Assignment{
      _functionDefinition.debugData,
      {Identifier{_functionDefinition.debugData, _functionDefinition.returnVariables.front().name}},
      std::make_unique<Expression>(generateMemoryLoad(
        m_context.dialect,
        _functionDefinition.debugData,
        *m_memoryOffsetTracker(_functionDefinition.returnVariables.front().name)
      ))
    });
    return;
  }

  if (!memoryVariableInits.empty())
    _functionDefinition.body.statements = std::move(memoryVariableInits) + std::move(_functionDefinition.body.statements);

  _functionDefinition.returnVariables = _functionDefinition.returnVariables | ranges::views::filter(
    std::not_fn(m_memoryOffsetTracker)
  ) | ranges::to<NameWithDebugDataList>;
}

void StackToMemoryMover::operator()(Block& _block)
{
  using OptionalStatements = std::optional<std::vector<Statement>>;

  auto rewriteAssignmentOrVariableDeclarationLeftHandSide = [this](
    auto& _stmt,
    auto& _lhsVars
  ) -> OptionalStatements {
    using StatementType = std::decay_t<decltype(_stmt)>;

    auto debugData = _stmt.debugData;
    if (_lhsVars.size() == 1)
    {
      if (auto offset = m_memoryOffsetTracker(_lhsVars.front().name))
        return generateMemoryStore(
          m_context.dialect,
          debugData,
          *offset,
          _stmt.value ? *std::move(_stmt.value) : Literal{debugData, LiteralKind::Number, LiteralValue(u256{0})}
        );
      else
        return {};
    }
    std::vector<std::optional<LiteralValue>> rhsMemorySlots;
    if (_stmt.value)
    {
      FunctionCall const* functionCall = std::get_if<FunctionCall>(_stmt.value.get());
      yulAssert(functionCall, "");
      if (isBuiltinFunctionCall(*functionCall))
        rhsMemorySlots = std::vector<std::optional<LiteralValue>>(_lhsVars.size(), std::nullopt);
      else
      {
        yulAssert(std::holds_alternative<Identifier>(functionCall->functionName));
        rhsMemorySlots =
          m_functionReturnVariables.at(std::get<Identifier>(functionCall->functionName).name) |
          ranges::views::transform(m_memoryOffsetTracker) |
          ranges::to<std::vector<std::optional<LiteralValue>>>;
      }
    }
    else
      rhsMemorySlots = std::vector<std::optional<LiteralValue>>(_lhsVars.size(), std::nullopt);

    // Nothing to do, if the right-hand-side remains entirely on the stack and
    // none of the variables in the left-hand-side are moved.
    if (
      ranges::none_of(rhsMemorySlots, [](std::optional<LiteralValue> const& _slot) { return _slot.has_value(); }) &&
      !util::contains_if(_lhsVars, m_memoryOffsetTracker)
    )
      return {};

    std::vector<Statement> memoryAssignments;
    std::vector<Statement> variableAssignments;
    VariableDeclaration tempDecl{debugData, {}, std::move(_stmt.value)};

    yulAssert(rhsMemorySlots.size() == _lhsVars.size(), "");
    for (auto&& [lhsVar, rhsSlot]: ranges::views::zip(_lhsVars, rhsMemorySlots))
    {
      std::unique_ptr<Expression> rhs;
      if (rhsSlot)
        rhs = std::make_unique<Expression>(generateMemoryLoad(m_context.dialect, debugData, *rhsSlot));
      else
      {
        YulName tempVarName = m_nameDispenser.newName(lhsVar.name);
        tempDecl.variables.emplace_back(NameWithDebugData{lhsVar.debugData, tempVarName});
        rhs = std::make_unique<Expression>(Identifier{debugData, tempVarName});
      }

      if (auto offset = m_memoryOffsetTracker(lhsVar.name))
        memoryAssignments += generateMemoryStore(
          m_context.dialect,
          _stmt.debugData,
          *offset,
          std::move(*rhs)
        );
      else
        variableAssignments.emplace_back(StatementType{
          debugData,
          { std::move(lhsVar) },
          std::move(rhs)
        });
    }

    std::vector<Statement> result;
    if (tempDecl.variables.empty())
      result.emplace_back(ExpressionStatement{debugData, *std::move(tempDecl.value)});
    else
      result.emplace_back(std::move(tempDecl));
    reverse(memoryAssignments.begin(), memoryAssignments.end());
    result += std::move(memoryAssignments);
    reverse(variableAssignments.begin(), variableAssignments.end());
    result += std::move(variableAssignments);
    return OptionalStatements{std::move(result)};
  };

  util::iterateReplacing(
    _block.statements,
    [&](Statement& _statement) -> OptionalStatements
    {
      visit(_statement);
      if (auto* assignment = std::get_if<Assignment>(&_statement))
        return rewriteAssignmentOrVariableDeclarationLeftHandSide(*assignment, assignment->variableNames);
      else if (auto* varDecl = std::get_if<VariableDeclaration>(&_statement))
        return rewriteAssignmentOrVariableDeclarationLeftHandSide(*varDecl, varDecl->variables);
      return {};
    }
  );
}

void StackToMemoryMover::visit(Expression& _expression)
{
  ASTModifier::visit(_expression);
  if (Identifier* identifier = std::get_if<Identifier>(&_expression))
    if (auto offset = m_memoryOffsetTracker(identifier->name))
      _expression = generateMemoryLoad(m_context.dialect, identifier->debugData, *offset);
}

std::optional<LiteralValue> StackToMemoryMover::VariableMemoryOffsetTracker::operator()(YulName const& _variable) const
{
  if (m_memorySlots.count(_variable))
  {
    uint64_t slot = m_memorySlots.at(_variable);
    yulAssert(slot < m_numRequiredSlots, "");
    auto const memoryOffset = m_reservedMemory + 32 * (m_numRequiredSlots - slot - 1);
    return valueOfNumberLiteral(toCompactHexWithPrefix(memoryOffset));
  }
  else
    return std::nullopt;
}

std::optional<LiteralValue> StackToMemoryMover::VariableMemoryOffsetTracker::operator()(NameWithDebugData const& _variable) const
{
  return (*this)(_variable.name);
}

std::optional<LiteralValue> StackToMemoryMover::VariableMemoryOffsetTracker::operator()(Identifier const& _variable) const
{
  return (*this)(_variable.name);
}

Coverage Report

Created: 2026-09-14 06:38