Lower signatures exposing struct registers to byval struct pointers

lib/Transforms/NaCl/SimplifyStructRegSignatures.cpp

Index: lib/Transforms/NaCl/SimplifyStructRegSignatures.cpp
diff --git a/lib/Transforms/NaCl/SimplifyStructRegSignatures.cpp b/lib/Transforms/NaCl/SimplifyStructRegSignatures.cpp
new file mode 100644
index 0000000000000000000000000000000000000000..d004fe9e32e35fec6df3ed0b985c9b421b589e0f
--- /dev/null
+++ b/lib/Transforms/NaCl/SimplifyStructRegSignatures.cpp
@@ -0,0 +1,577 @@
+//===- SimplifyStructRegSignatures.cpp - Change struct regs to struct ----===//
+// pointers

[JF, 2015/03/12 18:36:18]

That's a bit weird formatting :-)

It would fit if you drop "Change"?

[Mircea Trofin, 2015/03/13 22:04:52]

Done.

+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+//
+// This pass replaces function signatures exposing struct registers
+// to byval pointer-based signatures.
+//
+// There are 2 types of signatures that are thus changed:
+//
+// @foo(%some_struct %val) -> @foo(%some_struct* byval %val)
+//      and
+// %someStruct @bar(<other_args>) -> void @bar(%someStruct* sret, <other_args>)
+//
+// Such function types may appear in other type declarations, for example:
+//
+// %a_struct = type { void (%some_struct)*, i32 }
+// 
+// We map such types to corresponding types, mapping the function types 
+// appropriately:
+//
+// %a_struct.0 = type { void (%some_struct*)*, i32 }
+//===----------------------------------------------------------------------===//
+
+#include <cassert>
+#include <cstddef>
+
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/DenseSet.h"
+#include "llvm/ADT/ilist.h"
+#include "llvm/ADT/SetVector.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/Twine.h"
+#include "llvm/IR/Argument.h"
+#include "llvm/IR/Attributes.h"
+#include "llvm/IR/BasicBlock.h"
+#include "llvm/IR/DerivedTypes.h"
+#include "llvm/IR/Function.h"
+#include "llvm/IR/GlobalValue.h"
+#include "llvm/IR/Instructions.h"
+#include "llvm/IR/Module.h"
+#include "llvm/IR/Type.h"
+#include "llvm/IR/Use.h"
+#include "llvm/IR/User.h"
+#include "llvm/IR/Value.h"
+#include "llvm/Pass.h"
+#include "llvm/PassInfo.h"
+#include "llvm/PassRegistry.h"
+#include "llvm/PassSupport.h"
+#include "llvm/Transforms/NaCl.h"
+#include "llvm/Support/Debug.h"
+
+using namespace llvm;
+
+namespace {
+class MappingResult {
+public:
+  MappingResult(Type *ATy, bool Chg) {
+    Ty = ATy;
+    Changed = Chg;
+  }
+
+  bool isChanged() { return Changed; }
+
+  Type *operator->() { return Ty; }
+
+  operator Type *() { return Ty; }
+
+private:
+  Type *Ty;
+  bool Changed;
+};
+
+// Utility class. For any given type, get the associated type that is free of
+// struct register arguments.
+class TypeMapper {
+public:
+  Type *getSimpleType(Type *Ty);
+
+private:
+  DenseMap<Type *, Type *> MappedTypes;
+  MappingResult getSimpleArgumentType(Type *Ty);
+  MappingResult getSimpleAggregateTypeInternal(Type *Ty);
+};
+
+// This is a ModulePass because the pass recreates functions in
+// order to change their signatures.
+class SimplifyStructRegSignatures : public ModulePass {
+public:
+  static char ID;
+
+  SimplifyStructRegSignatures() : ModulePass(ID) {
+    initializeSimplifyStructRegSignaturesPass(*PassRegistry::getPassRegistry());
+  }
+  virtual bool runOnModule(Module &M);
+
+private:
+  TypeMapper Mapper;
+  DenseSet<Function *> FunctionsToDelete;
+  SetVector<CallInst*> CallsToPatch;
+  SetVector<InvokeInst *> InvokesToPatch;
+  DenseMap<Function *, Function *> FunctionMap;
+  bool simplifyFunction(Function *OldFunc, Module &M);
+  void scheduleCallsForCleanup(Function *NewFunc);
+  template <class TCall> void fixCallSite(TCall *Call);
+  void fixFunctionBody(Function *OldFunc, Function *NewFunc);
+};
+}
+
+static const unsigned int TypicalFuncArity = 8;
+static const unsigned int TypicalStructArity = 8;
+
+char SimplifyStructRegSignatures::ID = 0;
+
+INITIALIZE_PASS(
+    SimplifyStructRegSignatures, "simplify-struct-reg-signatures",
+    "Simplify function signatures by removing struct register parameters",
+    false, false)
+
+// The type is "simple" if it does not recursively reference a
+// function type with at least an operand (arg or return) typed as struct
+// register
+Type *TypeMapper::getSimpleType(Type *Ty) {
+  if (MappedTypes.count(Ty))
+    return MappedTypes[Ty];

[JF, 2015/03/12 18:36:18]

auto Found = MappedTypes.find(Ty);
if (Found != MappedTypes.end())
  return *Found;

[Mircea Trofin, 2015/03/13 22:04:52]

or the location alternative, to avoid the double search :)

[JF, 2015/03/14 18:42:57]

There's no point in having two separate functions now, I'd just drop "Internal".

+  return MappedTypes[Ty] = getSimpleAggregateTypeInternal(Ty);
+}
+
+// transforms any type that could transitively reference a function pointer

[JF, 2015/03/12 18:36:18]

"Transforms"

[Mircea Trofin, 2015/03/13 22:04:51]

Done.

+// into a simplified type.
+MappingResult TypeMapper::getSimpleAggregateTypeInternal(Type *Ty) {
+  LLVMContext &Ctx = Ty->getContext();
+

[JF, 2015/03/12 18:36:18]

assert(MappedTypes.end() == MappedTypes.find(Ty));

Is this an invariant? External code accesses this function from getSimpleType so
it'll always check the map to make sure the types hasn't been rewritten yet, but
I can't convince myself that recursive calls keep the invariant (and I think
they should).

[Mircea Trofin, 2015/03/13 22:04:51]

They don't. Struct mapping will insert the empty new struct there. Also we may
hit memoized types when following pointers, for example. Alternatives would
include passing a working map around, but I'm not sure how much more
maintainable that is. I'll mark a TODO to find a simpler (more maintainable)
algorithm, if possible, as a separate task.

+  if (auto *OldFnTy = dyn_cast<FunctionType>(Ty)) {
+    Type *OldRetType = OldFnTy->getReturnType();
+    Type *NewRetType = OldRetType;
+    Type *Void = Type::getVoidTy(Ctx);
+    SmallVector<Type *, TypicalFuncArity> NewArgs;
+    bool HasChanges = false;

[JF, 2015/03/12 18:36:18]

"Changed" is a more common name in LLVM.

[Mircea Trofin, 2015/03/13 22:04:51]

Done.

+    // struct register returns become the first parameter of the new FT.
+    // the new FT has void for the return type
+    if (OldRetType->isAggregateType()) {
+      NewRetType = Void;
+      HasChanges = true;
+      NewArgs.push_back(getSimpleArgumentType(OldRetType));
+    }
+    for (auto OldParam = OldFnTy->param_begin(), E = OldFnTy->param_end();
+         OldParam != E; ++OldParam) {

[JF, 2015/03/12 18:36:17]

for (auto P : OldFnTy->params())

[Mircea Trofin, 2015/03/13 22:04:52]

Done.

+      auto NewType = getSimpleArgumentType(*OldParam);
+      HasChanges |= NewType.isChanged();
+      NewArgs.push_back(NewType);
+    }
+    Type *NewFuncType = FunctionType::get(NewRetType, NewArgs, false);

[JF, 2015/03/12 18:36:17]

This should preserve vararg (and have a corresponding test).

[Mircea Trofin, 2015/03/13 22:04:51]

Done.

+    return MappingResult(NewFuncType, HasChanges);

[JF, 2015/03/12 18:36:17]

C++11 initializer lists should work here and a few places below:
  s/return MappingResult(\(.*\));/return {\1};/g

[Mircea Trofin, 2015/03/13 22:04:51]

Done.

+  }
+
+  if (auto PtrTy = dyn_cast<PointerType>(Ty)) {
+    auto NewTy = getSimpleAggregateTypeInternal(PtrTy->getPointerElementType());
+
+    return MappingResult(NewTy->getPointerTo(), NewTy.isChanged());

[JF, 2015/03/12 18:36:17]

This should also preserve the address space of the pointer.

[Mircea Trofin, 2015/03/13 22:04:51]

Done.

+  }
+
+  if (auto ArrTy = dyn_cast<ArrayType>(Ty)) {
+    auto NewTy = getSimpleAggregateTypeInternal(ArrTy->getArrayElementType());
+    return MappingResult(ArrayType::get(NewTy, ArrTy->getArrayNumElements()),
+                         NewTy.isChanged());
+  }
+
+  if (auto VecTy = dyn_cast<VectorType>(Ty)) {
+    auto NewTy = getSimpleAggregateTypeInternal(VecTy->getVectorElementType());
+    return MappingResult(VectorType::get(NewTy, VecTy->getVectorNumElements()),
+                         NewTy.isChanged());

[JF, 2015/03/12 18:36:17]

There's no test for this.

+  }
+
+  if (auto StructTy = dyn_cast<StructType>(Ty)) {
+    if (!StructTy->isLiteral()) {
+      // LLVM doesn't intern identified structs (the ones with a name). This,
+      // together with the fact that such structs can be recursive,
+      // complicates things a bit. We want to make sure that we only change
+      // "unsimplified" structs (those that somehow reference funcs that
+      // are not simple).
+      // We don't want to change "simplified" structs, otherwise converting
+      // instruction types will become trickier.
+      Type *&Loc = MappedTypes[StructTy];

[JF, 2015/03/12 18:36:17]

I'd use `find` here too.

[Mircea Trofin, 2015/03/13 22:04:52]

Done.

+      if (!Loc) {
+        // We don't have a mapping, and we don't know if the struct is recursive
+        // so we create an empty one and hypothesize that it is the
+        // mapping.
+        Loc = StructType::create(Ctx, StructTy->getStructName());

[JF, 2015/03/12 18:36:16]

This loses `isPacked`, though I'm not sure we care?

[Mircea Trofin, 2015/03/13 22:04:51]

No, we're setting that further below, if the new struct "wins". see line 230

[Mircea Trofin, 2015/03/13 22:04:52]

It doesn't, we set it below, line 229

+      } else {
+        // We either have a finished mapping or this is the empty placeholder
+        // created above, and we are in the process of finalizing it.
+        // 1) if this is a mapping, it must have the same element count
+        // as the original struct, so we mark a change if the types are
+        // different objects
+        // 2) if this is a placeholder, the element count
+        // wgetStructNumElementsill differ.

[JF, 2015/03/12 18:36:17]

Paste-o?

[Mircea Trofin, 2015/03/13 22:04:51]

yup

+        // Since we don't know yet if this is a change or not - because we
+        // are constructing the mapping - we don't mark as change. We decide
+        // if it is a change below, based on the other struct elements.
+        bool hasChanged =
+            StructTy != Loc &&
+            StructTy->getStructNumElements() == Loc->getStructNumElements();
+        return MappingResult(Loc, hasChanged);

[JF, 2015/03/12 18:36:17]

Same as above, just "Changed".

+      }
+    }
+
+    SmallVector<Type *, TypicalStructArity> ElemTypes;
+    bool HasChanges = false;

[JF, 2015/03/12 18:36:17]

Changed

[Mircea Trofin, 2015/03/13 22:04:52]

Done.

+    unsigned StructElemCount = StructTy->getStructNumElements();
+    for (unsigned I = 0; I < StructElemCount; I++) {
+      auto NewElem =
+          getSimpleAggregateTypeInternal(Ty->getStructElementType(I));
+      ElemTypes.push_back(NewElem);
+      HasChanges |= NewElem.isChanged();
+    }
+    if (!HasChanges) {
+      // We are leaking the created struct here, but there is no way to
+      // correctly delete it.
+      return MappingResult(MappedTypes[Ty] = Ty, false);
+    }
+
+    if (StructTy->isLiteral()) {
+      return MappingResult(MappedTypes[Ty] = StructType::get(
+                               Ctx, ElemTypes, StructTy->isPacked()),
+                           HasChanges);
+    } else {
+      Type *&Loc = MappedTypes[StructTy];

[JF, 2015/03/12 18:36:18]

`find`

[Mircea Trofin, 2015/03/13 22:04:52]

Done.

+      assert(Loc);
+      StructType *NewStruct = dyn_cast<StructType>(Loc);

[JF, 2015/03/12 18:36:18]

Just `cast<StructType>`, which has an assert internally.

[Mircea Trofin, 2015/03/13 22:04:51]

Done.

+      NewStruct->setBody(ElemTypes, StructTy->isPacked());
+      return MappingResult(MappedTypes[Ty] = NewStruct, true);
+    }
+  }
+
+  // anything else stays the same.

[JF, 2015/03/12 18:36:18]

Capitalize (here and below).

+  return MappingResult(Ty, false);
+}
+
+// get the simplified type of a function argument.
+MappingResult TypeMapper::getSimpleArgumentType(Type *Ty) {
+  // struct registers become pointers to simple structs
+  if (Ty->isAggregateType()) {
+    return MappingResult(
+        PointerType::get(getSimpleAggregateTypeInternal(Ty), 0), true);
+  }
+
+  return getSimpleAggregateTypeInternal(Ty);
+}
+
+// apply 'byval' to func arguments that used to be struct regs.
+// apply 'sret' to the argument corresponding to the return in the old signature
+static void ApplyByValAndSRet(Function *OldFunc, Function *NewFunc) {
+  auto const &OldArgList = OldFunc->getArgumentList();
+  auto &NewArgList = NewFunc->getArgumentList();
+
+  // when calling addAttribute, the first one refers to the function, so we
+  // skip past that.
+  unsigned ArgOffset = 1;
+  if (OldFunc->getReturnType()->isAggregateType()) {
+    NewFunc->addAttribute(1, Attribute::AttrKind::StructRet);
+    ArgOffset++;
+  }
+
+  auto NewArg = NewArgList.begin();
+  for (const Argument &OldArg : OldArgList) {

[JF, 2015/03/12 18:36:18]

Move `OldFunc->getArgumentList()` here, delete above.

[Mircea Trofin, 2015/03/13 22:04:51]

Done.

+    if (OldArg.getType()->isAggregateType()) {
+      NewFunc->addAttribute(NewArg->getArgNo() + ArgOffset,
+                            Attribute::AttrKind::ByVal);
+    }
+    NewArg++;
+  }
+}
+
+// update the arg names for a newly created function
+static void UpdateArgNames(Function *OldFunc, Function *NewFunc) {
+  auto NewArgIter = NewFunc->arg_begin();
+  auto const &OldFuncArgs = OldFunc->args();
+  if (OldFunc->getReturnType()->isAggregateType()) {
+    NewArgIter->setName("retVal");
+    NewArgIter++;
+  }
+
+  for (const Argument &OldArg : OldFuncArgs) {

[JF, 2015/03/12 18:36:18]

Same, move `NewFunc->args()` here.

[Mircea Trofin, 2015/03/13 22:04:51]

Done.

+    Argument *NewArg = NewArgIter++;
+    if (OldArg.getType()->isAggregateType()) {
+      NewArg->setName(OldArg.getName() + ".ptr");
+    } else {
+      NewArg->setName(OldArg.getName());
+    }
+  }
+}
+
+// replace all uses of an old value with a new one, disregarding the type. We
+// correct the types separately

[JF, 2015/03/12 18:36:17]

Where is the type corrected? It's not obvious from the surrounding lines and
below usage of BlindReplace.

[Mircea Trofin, 2015/03/13 22:04:52]

Done.

+static void BlindReplace(Value *Old, Value *New) {
+  for (auto UseIter = Old->use_begin(), E = Old->use_end(); E != UseIter;) {

[JF, 2015/03/12 18:36:17]

`for (auto U : Old->uses())`

Actually, it would be better to use the approach RAUW has of `while
(!use_empty())`.

[Mircea Trofin, 2015/03/13 22:04:52]

That wouldn't work, I'm not changing the content of the use list, just the uses.
It's not exactly RAUW, for that reason

+    Use &AUse = *(UseIter++);
+    AUse.set(New);
+  }
+}
+
+// adapt the body of a function for the new arguments
+static void ConvertArgumentValue(Value *Old, Value *New,
+                                 Instruction *InsPoint) {
+  if (Old == New)
+    return;
+
+  if (Old->getType() == New->getType()) {
+    Old->replaceAllUsesWith(New);
+    New->takeName(Old);
+    return;
+  }
+
+  if (Old->getType()->isAggregateType() && New->getType()->isPointerTy()) {
+    Value *Load = new LoadInst(New, Old->getName() + ".sreg", InsPoint);
+    BlindReplace(Old, Load);
+  } else {
+    BlindReplace(Old, New);
+  }
+}
+
+// fix returns. Return true if fixes were needed
+static void FixReturn(Function *OldFunc, Function *NewFunc) {
+
+  Argument *FirstNewArg = NewFunc->getArgumentList().begin();
+
+  for (auto BIter = NewFunc->begin(), LastBlock = NewFunc->end();
+       LastBlock != BIter;) {

[JF, 2015/03/12 18:36:19]

for (BasicBlock &BB : NewFunc->getBasicBlockList())

[Mircea Trofin, 2015/03/13 22:04:52]

no, because we mutate the list (same goes for the rest of similar comments). 

+    BasicBlock *BB = BIter++;
+    for (auto IIter = BB->begin(), LastI = BB->end(); LastI != IIter;) {

[JF, 2015/03/12 18:36:17]

for (Instruction &I : BB)

+      Instruction *Instr = IIter++;
+      if (ReturnInst *Ret = dyn_cast<ReturnInst>(Instr)) {
+        auto &Ctx = Ret->getContext();

[JF, 2015/03/12 18:36:18]

This should be loop invariant, hoist it.

[Mircea Trofin, 2015/03/13 22:04:52]

Done.

+        auto RetVal = Ret->getReturnValue();
+        ReturnInst *NewRet = ReturnInst::Create(Ctx, nullptr, Ret);
+        Ret->eraseFromParent();

[JF, 2015/03/12 18:36:18]

I think this invalidates your Instruction iterator. You should create a worklist
first, and then replace+delete (or create a kill list).

[Mircea Trofin, 2015/03/13 22:04:52]

That's the reason I increment the iterator first thing - which gets me a valid
iterator at all times.

+        Ret = nullptr;
+        new StoreInst(RetVal, FirstNewArg, NewRet);

[JF, 2015/03/12 18:36:16]

Give it a name and copy debug location.

Also, you can guarantee preferred alignment if you also do the same in the
alloca.

[Mircea Trofin, 2015/03/13 22:04:51]

Turns out store doesn't accept a name. Otherwise, done.

+      }
+    }
+  }
+}
+
+template <class TCall>
+void CopyCallAttributesAndMetadata(TCall *Orig, TCall *NewCall) {
+  NewCall->setCallingConv(Orig->getCallingConv());
+  NewCall->setAttributes(NewCall->getAttributes().addAttributes(
+      Orig->getContext(), AttributeSet::FunctionIndex,
+      Orig->getAttributes().getFnAttributes()));
+  NewCall->setDebugLoc(Orig->getDebugLoc());

[JF, 2015/03/12 18:36:17]

Other metadata?

[Mircea Trofin, 2015/03/13 22:04:51]

I believe this is it, but I added a TODO, because we may want to centralize this
copying of metadata for instructions, and push it upstream. I'll investigate
this separately.

+}
+
+static InvokeInst *CreateCallFrom(InvokeInst *Orig, Value *Target,
+                                  ArrayRef<Value *> &Args) {
+  InvokeInst *Ret = InvokeInst::Create(Target, Orig->getNormalDest(),
+                                       Orig->getUnwindDest(), Args);

[JF, 2015/03/12 18:36:17]

Copy name.

[Mircea Trofin, 2015/03/13 22:04:52]

Done.

+  CopyCallAttributesAndMetadata(Orig, Ret);
+  return Ret;
+}
+
+static CallInst *CreateCallFrom(CallInst *Orig, Value *Target,
+                                ArrayRef<Value *> &Args) {
+
+  CallInst *Ret = CallInst::Create(Target, Args);

[JF, 2015/03/12 18:36:17]

Name.

[Mircea Trofin, 2015/03/13 22:04:52]

Done, doing this in CopyCallAttributesAndMetadata

+
+  Ret->setTailCallKind(Orig->getTailCallKind());
+  CopyCallAttributesAndMetadata(Orig, Ret);

[JF, 2015/03/12 18:36:18]

I think you're missing NoBuiltin, NoInline, ReturnTwice, NotAccessMemory,
OnlyReadsMemory, NotReturn, NotThrow, CannotDuplicate.

[Mircea Trofin, 2015/03/13 22:04:52]

as per http://llvm.org/docs/LangRef.html#i-call:

"8. The optional function attributes list. Only ‘noreturn‘, ‘nounwind‘,
‘readonly‘ and ‘readnone‘ attributes are valid here."

+  return Ret;
+}
+
+// fix a call site by handing return type changes and/or parameter type and
+// attribute changes
+template <class TCall>
+void SimplifyStructRegSignatures::fixCallSite(TCall *OldCall) {
+  Value *NewTarget = OldCall->getCalledValue();
+
+  if (Function *CalledFunc = dyn_cast<Function>(NewTarget)) {
+    NewTarget = this->FunctionMap[CalledFunc];
+  }
+  assert(NewTarget);
+
+  FunctionType *NewType = dyn_cast<FunctionType>(

[JF, 2015/03/12 18:36:16]

cast<FunctionType>

[Mircea Trofin, 2015/03/13 22:04:51]

Done.

+      Mapper.getSimpleType(NewTarget->getType())->getPointerElementType());
+
+  Type *OldRetType = OldCall->getType();
+  const bool isSRet =
+      !OldCall->getType()->isVoidTy() && NewType->getReturnType()->isVoidTy();
+
+  const unsigned argOffset = isSRet ? 1 : 0;
+
+  SmallVector<Value *, TypicalFuncArity> NewArgs;
+
+  if (isSRet) {
+    AllocaInst *Alloca = new AllocaInst(OldRetType);

[JF, 2015/03/12 18:36:18]

Alignment to preferred size.

Name.

You can also pass OldCall instead of doing insertBefore below.

[Mircea Trofin, 2015/03/13 22:04:52]

It's simpler to use takeName than cache the old name and call the AllocaInst
constructor with the name and Instruction *InsertBefore signature.

+    NewArgs.push_back(Alloca);
+    Alloca->insertBefore(OldCall);
+
+    LoadInst *Load = new LoadInst(Alloca, OldCall->getName() + ".sreg",
+                                  (Instruction *)nullptr);

[JF, 2015/03/12 18:36:18]

Alignment.

[Mircea Trofin, 2015/03/13 22:04:52]

Done.

+    Load->insertAfter(OldCall);
+    OldCall->replaceAllUsesWith(Load);
+  }

[JF, 2015/03/12 18:36:18]

Why not create the new call between the alloca and the load? You'd be
constructing in a more intuitive order, and could use a straight IRBuilder.

[Mircea Trofin, 2015/03/13 22:04:52]

the call is between the alloca and the load *if* we have the return refactoring,
otherwise there's nothing to "in between". Or maybe I don't understand your
feedback?

[JF, 2015/03/14 18:42:57]

I mean that the order of code in this pass should mimic the order that code can
get generated in.

+
+  SmallSetVector<unsigned, TypicalFuncArity> ByRefPlaces;
+
+  for (unsigned ArgPos = 0;
+       ArgPos < NewType->getFunctionNumParams() - argOffset; ArgPos++) {
+
+    Use &OldArgUse = OldCall->getOperandUse(ArgPos);
+    Value *OldArg = OldArgUse;
+    Type *OldArgType = OldArg->getType();
+    unsigned NewArgPos = OldArgUse.getOperandNo() + argOffset;
+    Type *NewArgType = NewType->getFunctionParamType(NewArgPos);
+
+    if (OldArgType != NewArgType && OldArgType->isAggregateType()) {
+      AllocaInst *Alloca =
+          new AllocaInst(OldArgType, OldArg->getName() + ".ptr", OldCall);
+      new StoreInst(OldArg, Alloca, OldCall);
+      ByRefPlaces.insert(NewArgPos);
+      NewArgs.push_back(Alloca);
+    } else {
+      NewArgs.push_back(OldArg);
+    }
+  }
+
+  ArrayRef<Value *> ArrRef = NewArgs;
+  TCall *NewCall = CreateCallFrom(OldCall, NewTarget, ArrRef);
+
+  // copy the attributes over, and add byref/sret as necessary
+  const AttributeSet &OldAttrSet = OldCall->getAttributes();
+  const AttributeSet &NewAttrSet = NewCall->getAttributes();
+  LLVMContext &Ctx = OldCall->getContext();
+  AttrBuilder Builder(OldAttrSet, 0);
+
+  for (unsigned I = 0; I < NewCall->getNumArgOperands(); I++) {
+    NewCall->setAttributes(NewAttrSet.addAttributes(
+        Ctx, I + argOffset + 1, OldAttrSet.getParamAttributes(I + 1)));
+    if (ByRefPlaces.count(I)) {
+      NewCall->addAttribute(I + 1, Attribute::AttrKind::ByVal);
+    }
+  }
+
+  if (isSRet) {
+    NewAttrSet.addAttributes(Ctx, 1, OldAttrSet.getRetAttributes());
+    NewCall->addAttribute(1, Attribute::AttrKind::StructRet);
+  } else {
+    NewCall->setAttributes(NewAttrSet.addAttributes(
+        Ctx, AttributeSet::ReturnIndex, OldAttrSet.getRetAttributes()));
+    // if we still return something, this is the value to replace the old
+    // call with
+    OldCall->replaceAllUsesWith(NewCall);
+  }
+
+  NewCall->insertBefore(OldCall);
+  OldCall->eraseFromParent();
+  OldCall = NULL;
+}
+
+void SimplifyStructRegSignatures::scheduleCallsForCleanup(Function *NewFunc) {
+  for (auto &BBIter : NewFunc->getBasicBlockList()) {
+    for (auto &IIter : BBIter.getInstList()) {
+      if (CallInst *Call = dyn_cast<CallInst>(&IIter)) {
+        CallsToPatch.insert(Call);
+      } else if (InvokeInst *Invoke = dyn_cast<InvokeInst>(&IIter)) {
+        InvokesToPatch.insert(Invoke);
+      }
+    }
+  }
+}
+
+// change function body in the light of type changes
+void SimplifyStructRegSignatures::fixFunctionBody(Function *OldFunc,
+                                                  Function *NewFunc) {
+  if (NewFunc->empty())
+    return;
+
+  bool returnWasFixed = OldFunc->getReturnType()->isAggregateType();
+
+  Instruction *InsPoint = NewFunc->begin()->begin();
+  auto NewArgIter = NewFunc->arg_begin();
+  // advance one more if we used to return a struct register
+  if (returnWasFixed)
+    NewArgIter++;
+
+  // wire new parameters in
+  for (auto ArgIter = OldFunc->arg_begin(), E = OldFunc->arg_end();
+       E != ArgIter;) {
+    Argument *OldArg = ArgIter++;
+    Argument *NewArg = NewArgIter++;
+    ConvertArgumentValue(OldArg, NewArg, InsPoint);
+  }
+
+  // now fix instruction types. Calls are dealt with separately, but we still
+  // update the types here. We know that each value could only possibly be
+  // of a simplified type. At the end of this, call sites will be invalid, but
+  // we handle that afterwards, to make sure we have all the functions changed
+  // first (so that calls have valid targets)
+  for (auto BBIter = NewFunc->begin(), LBlock = NewFunc->end();
+       LBlock != BBIter;) {
+    auto Block = BBIter++;
+    for (auto IIter = Block->begin(), LIns = Block->end(); LIns != IIter;) {
+      auto Instr = IIter++;
+      Instr->mutateType(Mapper.getSimpleType(Instr->getType()));
+    }
+  }
+  if (returnWasFixed)
+    FixReturn(OldFunc, NewFunc);
+}
+
+// Ensure function is simplified, returning true if the function
+// had to be changed.
+bool SimplifyStructRegSignatures::simplifyFunction(Function *OldFunc,
+                                                   Module &M) {
+  FunctionType *OldFT = OldFunc->getFunctionType();
+  FunctionType *NewFT = dyn_cast<FunctionType>(Mapper.getSimpleType(OldFT));
+  assert(NewFT);
+
+  Function *&AssociatedFctLoc = FunctionMap[OldFunc];
+  if (NewFT != OldFT) {
+    Function *NewFunc = Function::Create(NewFT, OldFunc->getLinkage());
+    AssociatedFctLoc = NewFunc;
+
+    NewFunc->copyAttributesFrom(OldFunc);
+    OldFunc->getParent()->getFunctionList().insert(OldFunc, NewFunc);
+    NewFunc->takeName(OldFunc);
+
+    UpdateArgNames(OldFunc, NewFunc);
+    ApplyByValAndSRet(OldFunc, NewFunc);
+
+    NewFunc->getBasicBlockList().splice(NewFunc->begin(),
+                                        OldFunc->getBasicBlockList());
+
+    fixFunctionBody(OldFunc, NewFunc);
+    FunctionsToDelete.insert(OldFunc);
+  } else {
+    AssociatedFctLoc = OldFunc;
+  }
+  scheduleCallsForCleanup(AssociatedFctLoc);
+  return NewFT != OldFT;
+}
+
+bool SimplifyStructRegSignatures::runOnModule(Module &M) {
+  bool Changed = false;
+
+  // change function signatures and fix a changed function body by
+  // wiring the new arguments. Call sites are unchanged at this point
+  for (Module::iterator Iter = M.begin(), E = M.end(); Iter != E;) {
+    Function *Func = Iter++;
+    Changed |= simplifyFunction(Func, M);
+  }
+
+  // fix call sites
+  for (auto &CallToFix : CallsToPatch) {
+    fixCallSite(CallToFix);
+  }
+
+  for (auto &InvokeToFix : InvokesToPatch) {
+    fixCallSite(InvokeToFix);
+  }
+
+  // delete leftover functions - the ones with old signatures
+  for (auto &ToDelete : FunctionsToDelete) {
+    // this also frees the memory
+    ToDelete->eraseFromParent();
+  }
+  return Changed;
+}
+
+ModulePass *llvm::createSimplifyStructRegSignaturesPass() {
+  return new SimplifyStructRegSignatures();
+}