DirectXShaderCompiler/lib/DxrFallback/FunctionBuilder.h

#pragma once

#include "llvm/IR/Module.h"

#include <string>
#include <vector>

//==============================================================================
// Simplifies the creation of functions.
//
// To create a function 'void foo( userType, i32, float* )' use the following
// code:
//   FunctionBuilder(module, "foo").voidTy().type(userType).i32().floatPtr().build()
//
// The first type specified is the return type.
class FunctionBuilder
{
public:
  FunctionBuilder(llvm::Module* mod, const std::string& name)
    : m_context(mod->getContext())
    , m_module(mod)
    , m_name(name)
  {}

  FunctionBuilder& voidTy()
  {
    m_argNames.push_back("");
    m_types.push_back(llvm::Type::getVoidTy(m_context));
    return *this;
  }
  FunctionBuilder& floatTy(const std::string& argName = "")
  {
    m_argNames.push_back(argName);
    m_types.push_back(llvm::Type::getFloatTy(m_context));
    return *this;
  }
  FunctionBuilder& floatPtr(const std::string& argName = "")
  {
    m_argNames.push_back(argName);
    m_types.push_back(llvm::Type::getFloatPtrTy(m_context));
    return *this;
  }
  FunctionBuilder& doubleTy(const std::string& argName = "")
  {
    m_argNames.push_back(argName);
    m_types.push_back(llvm::Type::getDoubleTy(m_context));
    return *this;
  }
  FunctionBuilder& doublePtr(const std::string& argName = "")
  {
    m_argNames.push_back(argName);
    m_types.push_back(llvm::Type::getDoublePtrTy(m_context));
    return *this;
  }
  FunctionBuilder& i32(const std::string& argName = "")
  {
    m_argNames.push_back(argName);
    m_types.push_back(llvm::Type::getInt32Ty(m_context));
    return *this;
  }
  FunctionBuilder& i32Ptr(const std::string& argName = "")
  {
    m_argNames.push_back(argName);
    m_types.push_back(llvm::Type::getInt32PtrTy(m_context));
    return *this;
  }
  FunctionBuilder& i16(const std::string& argName = "")
  {
    m_argNames.push_back(argName);
    m_types.push_back(llvm::Type::getInt16Ty(m_context));
    return *this;
  }
  FunctionBuilder& i16Ptr(const std::string& argName = "")
  {
    m_argNames.push_back(argName);
    m_types.push_back(llvm::Type::getInt16PtrTy(m_context));
    return *this;
  }
  FunctionBuilder& i8(const std::string& argName = "")
  {
    m_argNames.push_back(argName);
    m_types.push_back(llvm::Type::getInt8Ty(m_context));
    return *this;
  }
  FunctionBuilder& i8Ptr(const std::string& argName = "")
  {
    m_argNames.push_back(argName);
    m_types.push_back(llvm::Type::getInt8PtrTy(m_context));
    return *this;
  }
  FunctionBuilder& i1(const std::string& argName = "")
  {
    m_argNames.push_back(argName);
    m_types.push_back(llvm::Type::getInt1Ty(m_context));
    return *this;
  }
  FunctionBuilder& i1Ptr(const std::string& argName = "")
  {
    m_argNames.push_back(argName);
    m_types.push_back(llvm::Type::getInt1PtrTy(m_context));
    return *this;
  }

  FunctionBuilder& type(llvm::Type* ty, const std::string& argName = "")
  {
    m_argNames.push_back(argName);
    m_types.push_back(ty);
    return *this;
  }
  FunctionBuilder& types(const std::vector<llvm::Type*>& ty, const std::vector<std::string>& argNames)
  {
    if (argNames.empty())
      for (size_t i = 0; i < ty.size(); ++i)
        m_argNames.push_back("");
    m_types.insert(m_types.end(), ty.begin(), ty.end());
    return *this;
  }

  llvm::Function* build()
  {
    using namespace llvm;

    Type*        retTy = m_types[0];
    AttributeSet attributes;
    Type**       argsBegin = (&m_types[0]) + 1;
    Type**       argsEnd = argsBegin + m_types.size() - 1;
    Constant*    funcC =
      m_module->getOrInsertFunction(m_name, FunctionType::get(retTy, ArrayRef<Type*>(argsBegin, argsEnd), false), attributes);
    Function* func = cast<Function>(funcC);

    std::string* argNamePtr = m_argNames.data() + 1;
    for (auto& arg : func->args())
      arg.setName(*argNamePtr++);

    return func;
  }

private:
  llvm::LLVMContext&       m_context;
  llvm::Module*            m_module = nullptr;
  std::string              m_name;
  std::vector<std::string> m_argNames;
  std::vector<llvm::Type*> m_types;

  // forbidden
  FunctionBuilder();
  FunctionBuilder(const FunctionBuilder&);
};