//===- LoopInterchange.cpp - Loop Interchange Pass ------------------------------===//
//
//                     The LLVM Compiler Infrastructure
//
// This file is distributed under the University of Illinois Open Source
// License. See LICENSE.TXT for details.
//
//===----------------------------------------------------------------------===//
//
// This file implements Loop Interchange Pass.
//
//===----------------------------------------------------------------------===//

#define DEBUG_TYPE "loop-interchange"
#include "llvm/Transforms/Scalar.h"
#include "llvm/Function.h"
#include "llvm/IntrinsicInst.h"
#include "llvm/Analysis/LoopPass.h"
#include "llvm/Analysis/Dominators.h"
#include "llvm/Analysis/ScalarEvolution.h"
#include "llvm/Transforms/Utils/Local.h"
#include "llvm/Transforms/Utils/BasicBlockUtils.h"
#include "llvm/Transforms/Utils/SSAUpdater.h"
#include "llvm/Support/CommandLine.h"
#include "llvm/Support/Debug.h"
#include "llvm/ADT/Statistic.h"
#include "llvm/ADT/SmallVector.h"
using namespace llvm;

#define MAX_HEADER_SIZE 16
STATISTIC(NumInterchanged, "Number of loops interchanged");
namespace {

  class LoopInterchange : public LoopPass {
  public:
    static char ID; // Pass ID, replacement for typeid
    LoopInterchange() : LoopPass(&ID)     {}

    // Interchange Loop L as many times as possible. Return true if
    // loop is interchanged at least once.
    bool runOnLoop(Loop *L, LPPassManager &LPM);

    // LCSSA form makes instruction renaming easier.
    virtual void getAnalysisUsage(AnalysisUsage &AU) const {
      AU.addRequiredID(LoopSimplifyID);
      AU.addPreservedID(LoopSimplifyID);
      AU.addRequiredID(LCSSAID);
      AU.addPreservedID(LCSSAID);      
      AU.addPreserved<ScalarEvolution>();
      AU.addRequired<LoopInfo>();
      AU.addPreserved<LoopInfo>();
      //      AU.addPreserved<DominatorTree>();
      //      AU.addPreserved<DominanceFrontier>();
    }

    // Helper functions

    /// Do actual work
    // bool interchangeLoop(Loop *L, LPPassManager &LPM5);
    

    /// Initialize local data
    void initialize(Loop* Lp);

  private:
    bool tightlyNested(Loop* inner, Loop* outer);
    //    void fixPHINodesAfterSplit(BasicBlock*, BasicBlock*);
    void restructureLoops(Loop* innerLoop, Loop* outerLoop);
    void splitInnerLoopLatch(DominatorTree*, LoopInfo&, Instruction*);
    void splitInnerLoopHeader(DominatorTree*, LoopInfo&);
    void updateDominators(DominatorTree* DT, BasicBlock* oldDominator, BasicBlock* newDominator);
    bool loopPairInterchanged(Loop* inner, Loop* outer);
    void updateSSA();
    void updateSSAagain();
    bool tripCountIsLoopInvariant(Loop* loop);
    Loop* innerLoop;
    Loop* outerLoop;
    BasicBlock *innerLoopHeader;
    BasicBlock *outerLoopHeader;
    BasicBlock *innerLoopLatch;
    BasicBlock *outerLoopLatch;
    BasicBlock* outerLoopPreHeader;
    PHINode* innerIndexVar;
    PHINode* outerIndexVar;
    //    Instruction* innerIndexVarInc;
    BasicBlock *innerLoopLatchPred;
    BasicBlock *innerLoopHeaderSucc;
    std::vector<std::pair<Loop*,Loop*> > interchangedLoops;
  
  };
}  
char LoopInterchange::ID = 0;
static RegisterPass<LoopInterchange> X("loop-interchange", "Interchange Loops");

Pass *llvm::createLoopInterchangePass() { return new LoopInterchange(); }

void LoopInterchange::initialize(Loop* Lp)
{
  innerLoop = Lp;
  outerLoop = innerLoop->getParentLoop();
  innerLoopHeader = innerLoop->getHeader();
  outerLoopHeader = outerLoop->getHeader();
  innerLoopLatch = innerLoop->getLoopLatch();
  outerLoopLatch = outerLoop->getLoopLatch();
  assert(innerLoopLatch && "Inner loop has multiple exit blocks");
  assert(outerLoopLatch && "Outer loop has multiple exit blocks");
  outerLoopPreHeader = outerLoop->getLoopPreheader();
  innerIndexVar = innerLoop->getCanonicalInductionVariable();
  outerIndexVar = outerLoop->getCanonicalInductionVariable();
  

}
bool LoopInterchange::tightlyNested(Loop* innerLoop, Loop* outerLoop)
{
  // Check if the outer loop has any basic blocks that are not in the inner loop.
  Loop::block_iterator iter_b = outerLoop->block_begin();
  Loop::block_iterator iter_end = outerLoop->block_end();
  for(; iter_b != iter_end; ++iter_b) {
    if (*iter_b != outerLoopHeader && *iter_b != outerLoopLatch &&
	innerLoop->contains(*iter_b) == false) { 
	  // The loops are not tightly nested. The interchange failed.
      return false;
    }
  }
  return true;
}
bool LoopInterchange::loopPairInterchanged(Loop* innerLoop, Loop* outerLoop)
{
  if ( std::find(interchangedLoops.begin(), 
	    interchangedLoops.end(), 
	    std::pair<Loop*, Loop*> (outerLoop, innerLoop)) !=
       interchangedLoops.end()) {
    return true; // This pair has been swapped already.
  }
  return false;
}

void LoopInterchange::updateDominators(DominatorTree* DT,BasicBlock* oldDominator, BasicBlock* newDominator)
{
  for (Loop::block_iterator BI = innerLoop->block_begin(), BE = innerLoop->block_end();
       BI != BE; ++BI) {
    BasicBlock *B = *BI;
    if ( B != oldDominator) {
      DomTreeNode *Node = DT->getNode(B);
      if (Node && Node->getBlock() == oldDominator)
	DT->changeImmediateDominator(*BI, newDominator);
    }
  }
}
void LoopInterchange::splitInnerLoopLatch(DominatorTree* DT, LoopInfo& LI, Instruction* inc)
{
  BasicBlock::iterator I = innerLoopLatch->begin();
  BasicBlock::iterator E = innerLoopLatch->end();
  for (;I != E; ++I) {
    if (inc == I) 
      break;
  }

  // Split the inner loop latch out.
  innerLoopLatchPred = innerLoopLatch; 
  innerLoopLatch = innerLoopLatch->splitBasicBlock(I);
  innerLoop->addBasicBlockToLoop(innerLoopLatch, LI.getBase());
  DT->addNewBlock(innerLoopLatch, innerLoopLatchPred);

  DT->changeImmediateDominator(outerLoopLatch, innerLoopLatchPred);
 
  updateDominators(DT, innerLoopLatchPred, innerLoopLatch);
}	

void LoopInterchange::splitInnerLoopHeader(DominatorTree* DT, LoopInfo& LI)
{

  // Split the inner loop header out.
  innerLoopHeaderSucc =  
    innerLoopHeader->splitBasicBlock(innerLoopHeader->getFirstNonPHI());
  DT->splitBlock(innerLoopHeaderSucc);
  DT->changeImmediateDominator(innerLoopHeaderSucc, innerLoopHeader);
  updateDominators(DT, innerLoopHeader,innerLoopHeaderSucc);
  outerLoop->addBasicBlockToLoop(innerLoopHeaderSucc, LI.getBase());
  if (!innerLoop->contains(innerLoopLatch))
    innerLoop->addBasicBlockToLoop(innerLoopHeaderSucc, LI.getBase());
  if (innerLoopLatchPred==innerLoopHeader)
    innerLoopLatchPred = innerLoopHeaderSucc;

  //  fixPHINodesAfterSplit(innerLoopHeader, innerLoopHeaderSucc);
}  
// Fix the PHINodes 
void LoopInterchange::restructureLoops(Loop* innerLoop, Loop* outerLoop)
{
  Loop* outerLoopParent = outerLoop->getParentLoop();
  LoopInfo &LI = getAnalysis<LoopInfo>();

  if (outerLoopParent) {
    // Remove the loop from its parent loop.
    for (Loop::iterator I = outerLoopParent->begin(), E = outerLoopParent->end();;
	 ++I) {
      assert(I != E && "Couldn't find loop");
      if (*I == outerLoop) {
	outerLoopParent->removeChildLoop(I);
	break;
      }
    }
    outerLoopParent->addChildLoop(innerLoop);
  } else {
    for (Loop::iterator I = outerLoop->begin(), E = outerLoop->end();;
	 ++I) {
      assert(I != E && "Couldn't find loop");
      if (*I == innerLoop) {
	outerLoop->removeChildLoop(I);
	break;
      }
    }
    LI.changeTopLevelLoop(outerLoop, innerLoop);
  }
  // Move child loop(s) from the inner loop to the outer loop before moving the
  // outer loop into the inner.
  for (Loop::iterator I = innerLoop->begin(), E = innerLoop->end(); 
       I != E; ++I) {
    outerLoop->addChildLoop(innerLoop->removeChildLoop(I));
  }
  
  // Move the body of the inner loop to the outer loop.
  {
    Loop::block_iterator I = innerLoop->block_begin();
    Loop::block_iterator E = innerLoop->block_end();
    for (; I != E; ++I) {
      if (*I != innerLoopHeader && *I != innerLoopLatch)
	if (!outerLoop->contains(*I))
	  outerLoop->addBlockEntry(*I);
    }
  }
  innerLoop->addChildLoop(outerLoop);
  // Make sure that the loops are consistent with respect to containment.
  {
    Loop::block_iterator iter = outerLoop->block_begin();
    Loop::block_iterator iter_end = outerLoop->block_end();
    for(;iter != iter_end; ++iter) {
      BasicBlock* BB = *iter;
      if (!innerLoop->contains(BB)) {
	innerLoop->addBlockEntry(BB);
      }
    }
  }

}
bool LoopInterchange::tripCountIsLoopInvariant(Loop* loop)
{
  assert(loop != NULL);
  Value* tripCount = loop->getTripCount();
  if (!tripCount)
    return false;
  return loop->isLoopInvariant(tripCount);
}
/// Interchange Loop L as many times as possible. Return true if
/// the loop is interchanged at least once.
bool LoopInterchange::runOnLoop(Loop *Lp, LPPassManager &LPM) {
  BasicBlock::iterator I,E;
  DominatorTree *DT = getAnalysisIfAvailable<DominatorTree>();
  LoopInfo &LI = getAnalysis<LoopInfo>();
  Instruction* innerIndexVarInc = Lp->getCanonicalInductionVariableIncrement();
  if (!Lp->getParentLoop())
    return false;
  initialize(Lp);

  // Check if the loop pair is already interchanged.
  if (loopPairInterchanged(innerLoop, outerLoop))
    return false;
  // Check if the loops are tightly nested.
  if (!tightlyNested(innerLoop, outerLoop))
    return false;

  // Check if the trip count is loop invariant
  if (!tripCountIsLoopInvariant(innerLoop) || 
      !tripCountIsLoopInvariant(outerLoop))
    return false;

  // Check if the loop is in LCSSA form
  if (!innerLoop->isLCSSAForm() || !outerLoop->isLCSSAForm())
    return false;

  // Check if the loop is in LCSSA form
  if (!innerLoop->isLoopSimplifyForm() || !outerLoop->isLoopSimplifyForm())
    return false;

  // Check if the loop is in normal form.
  if (!innerLoop->isLoopSimplifyForm() || !outerLoop->isLoopSimplifyForm())
    return false;

  // Check if the loops are in simplified form.
  if (!innerLoop->getCanonicalInductionVariable() || 
      !outerLoop->getCanonicalInductionVariable()) {
    return false;
  }

  // TODO Perform data dependence analysis for the loop nest.
  // TODO Check if there is any data dependence preventing interchange.

  updateSSA();

  restructureLoops(innerLoop, outerLoop);

  splitInnerLoopLatch(DT, LI, innerIndexVarInc);

  splitInnerLoopHeader(DT, LI);

  // Outer loop preheader transfers to the inner loop header
  BranchInst* BI = dyn_cast<BranchInst>(outerLoopPreHeader->getTerminator());
  if (!BI)
    return false;
  if (BI->isUnconditional()) {
    BI->setUnconditionalDest(innerLoopHeader);
  } else {
    if (BI->getSuccessor(0) == outerLoopHeader) {
      BI->setSuccessor(0, innerLoopHeader);
    } else {
      BI->setSuccessor(1, innerLoopHeader);
    }
  }
  
  
  // Inner loop header preceeds outer loop header.
  BI = dyn_cast<BranchInst> (innerLoopHeader->getTerminator());
  if (!BI)
    return false;
  BasicBlock* succ=NULL;
  if (BI->isUnconditional()) {
    succ = BI->getSuccessor(0);
    BI->setSuccessor(0, outerLoopHeader);
  } else {
    if (innerLoop->contains(BI->getSuccessor(0))) {
      succ = BI->getSuccessor(0);
      BI->setSuccessor(0, outerLoopHeader);
    } else if (innerLoop->contains(BI->getSuccessor(1))) {
      succ = BI->getSuccessor(1);
      BI->setSuccessor(1, outerLoopHeader);
    }
  }
  DT->changeImmediateDominator(outerLoopHeader, innerLoopHeader);
  
  // Outer loop header transfers to the successor of the inner loop header.
  BI = dyn_cast<BranchInst> (outerLoopHeader->getTerminator());
  if (BI->isUnconditional()) {
    BI->setSuccessor(0, succ);
  } else {
    if (outerLoop->contains(BI->getSuccessor(0))) {
      BI->setSuccessor(0, succ);
    } else if (outerLoop->contains(BI->getSuccessor(1))) {
      BI->setSuccessor(1, succ);
    }
  }
  DT->changeImmediateDominator(succ, outerLoopHeader);
  
  // Outer loop latch exits to the inner loop latch.
  BasicBlock* outerLoopSuccessor=NULL;
  BI = dyn_cast<BranchInst> (outerLoopLatch->getTerminator());
  if (BI->isUnconditional()) {    
    outerLoopSuccessor = BI->getSuccessor(0);
    BI->setSuccessor(0, innerLoopLatch);
  } else {
    if (outerLoop->contains(BI->getSuccessor(0))) {
      outerLoopSuccessor = BI->getSuccessor(1);
      BI->setSuccessor(1, innerLoopLatch);
    } else {
      outerLoopSuccessor = BI->getSuccessor(0);
      BI->setSuccessor(0, innerLoopLatch);
    }
  }
  
  // Inner loop latch exits out to the successor of the loop nest.
  BI = dyn_cast<BranchInst>  (innerLoopLatch->getTerminator());
  if (BI->isUnconditional()) {
    BI->setSuccessor(0, outerLoopSuccessor);
  }  else {
    if (BI->getSuccessor(0) == outerLoopLatch) {
      BI->setSuccessor(0, outerLoopSuccessor);
    } else
	BI->setSuccessor(1, outerLoopSuccessor);
  }
  DT->changeImmediateDominator(outerLoopSuccessor, innerLoopLatch);
  DT->changeImmediateDominator(innerLoopLatch, outerLoopLatch);
  
  BI = dyn_cast<BranchInst> (innerLoopLatchPred->getTerminator());
  BI->setSuccessor(0, outerLoopLatch);
  {
    Loop::block_iterator iter = innerLoop->block_begin();
    Loop::block_iterator iter_end = innerLoop->block_end();
    for(;iter != iter_end; ++iter) {
      if (*iter != innerLoopHeader && *iter != innerLoopLatch) {
	LI.changeLoopFor(*iter, outerLoop);
	if (!outerLoop->contains(*iter))
	  outerLoop->addBlockEntry(*iter);
      } else {
	LI.changeLoopFor(*iter, innerLoop);
      }
    }	
  }
  
  interchangedLoops.push_back(std::pair<Loop*, Loop*>((Loop*)innerLoop, (Loop*)outerLoop));
  
  
  
  // Fix the PHI nodes for the exit node of the outer loop by 
  // replacing the outer latch with inner latch.
  updateSSAagain();
  // Outer loop successor (outer loop exit node)
  for (BasicBlock::iterator I = outerLoopSuccessor->begin();
       PHINode* PN = dyn_cast<PHINode>(I); ++I) {
      int index = PN->getBasicBlockIndex(outerLoopLatch);
      if (index != -1)
	PN->setIncomingBlock(index, innerLoopLatch);
  }

  // innerLoopLatchPred feeds to outer loop latch instead of innerLoopLatch.
  for (BasicBlock::iterator I = outerLoopLatch->begin();
       PHINode* PN = dyn_cast<PHINode>(I); ++I) {
    int index = PN->getBasicBlockIndex(innerLoopLatch);
    if (index != -1)
      PN->setIncomingBlock(index, innerLoopLatchPred);
  }
  // Remove outer loop header -> inner loop header
  for (BasicBlock::iterator I = innerLoopHeader->begin();
       PHINode* PN = dyn_cast<PHINode>(I); ++I) {
      int index = PN->getBasicBlockIndex(outerLoopHeader);
      if (index != -1)
	PN->removeIncomingValue(index);
  }

  // Clone outer loop latch -> outer loop header into inner loop latch
  for (BasicBlock::iterator I = outerLoopHeader->begin();
       PHINode* PN = dyn_cast<PHINode>(I); ++I) {
      int index = PN->getBasicBlockIndex(outerLoopLatch);
      if (index != -1) {
	// Create a new assignment inst to carry to the next iteration
	// of the inner loop.
	// PN->clone()->insertBefore(innerLoopLatch->getTerminator());
      }
  }

  // Replace outer loop header -> inner loop header with 
  // inner loop latch -> inner loop header.
  for (BasicBlock::iterator I = outerLoopHeader->begin();
       PHINode* PN = dyn_cast<PHINode>(I); ++I) {
      int index = PN->getBasicBlockIndex(outerLoopHeader);
      if (index != -1) {
	PN->setIncomingBlock(index, innerLoopLatch);
      }
  }
  
  // Now update the outer loop header.
  //    outerLoop->removeBlockFromLoop(innerLoopHeaderSucc);
  NumInterchanged++;
  LPM.deleteLoopFromQueue(outerLoop); 
  return true;
}

void LoopInterchange::updateSSA()
{
  // Begin by walking outerLoopHeader and populating ValueMap with an entry for
  // each Instruction.
  BasicBlock::iterator I = outerLoopHeader->begin(), 
    E = outerLoopHeader->end();
  DenseMap<const Value *, Value *> ValueMap;
  for (; PHINode *PN = dyn_cast<PHINode>(I); ++I) {
    if (PN != outerIndexVar) {
      ValueMap[PN] = 
	PN->getIncomingValue(PN->getBasicBlockIndex(outerLoopPreHeader)); 
    } else {
      int index = PN->getBasicBlockIndex(outerLoopPreHeader);
      if (index != -1)
	PN->setIncomingBlock(index, innerLoopHeader);
    }
  }
  TerminatorInst *LoopEntryBranch = outerLoopPreHeader->getTerminator();
  for (; I != E; ++I) {
    Instruction *C = I->clone();
    C->setName(I->getName());
    C->insertBefore(LoopEntryBranch);
    ValueMap[I] = C;

  }
  LoopEntryBranch->eraseFromParent();

  for (I=innerLoopHeader->begin(); PHINode *PN = dyn_cast<PHINode>(I); ++I) 
    PN->addIncoming(PN->removeIncomingValue(outerLoopHeader), 
  		    outerLoopPreHeader);

  SSAUpdater SSA;
  for (I= outerLoopHeader->begin(); PHINode *PN = dyn_cast<PHINode>(I); ) {
    ++I;
    if (PN == outerIndexVar)
      continue;
    Value *outerLoopHeaderVal = PN;
    Value *outerLoopPreHeaderVal = ValueMap[PN];
    SSA.Initialize(PN);
    SSA.AddAvailableValue(outerLoopHeader, outerLoopHeaderVal);
    SSA.AddAvailableValue(outerLoopPreHeader, outerLoopPreHeaderVal);

    // Visit each use of the outerLoopHeader instruction.
    for (Value::use_iterator UI = outerLoopHeaderVal->use_begin(),
         UE = outerLoopHeaderVal->use_end(); UI != UE; ) {
      // Grab the use before incrementing the iterator.
      Use &use = UI.getUse();
      ++UI;
      Instruction *UserInst = cast<Instruction>(use.getUser());
      BasicBlock *UserBB = UserInst->getParent();
      if (UserBB == innerLoopHeader) {
	ValueMap[PN] = use.getUser();
      }
      if (!isa<PHINode>(UserInst)) {
        // Users in the outerLoopPreHeader need to use the value to which the
        // original definitions are mapped.
        if (UserBB == outerLoopPreHeader) {
          use = outerLoopPreHeaderVal;
          continue;
        }
      }
      // Anything else can be handled by SSAUpdater.
      SSA.RewriteUse(use);
    }
  }
  for (I=outerLoopHeader->begin(); PHINode *PN = dyn_cast<PHINode>(I); ++I) {
    if (PN == outerIndexVar) 
      continue;
    int index = PN->getBasicBlockIndex(outerLoopPreHeader);
    if (index != -1) {
      PN->setIncomingBlock(index, innerLoopHeader);
      PN->setIncomingValue(index, ValueMap[I]);
    }
  }
}
void LoopInterchange::updateSSAagain()
{
  DenseMap<const Value *, Value *> ValueMap;

  // Map incoming values from inner loop latch.
  for (BasicBlock::iterator I = innerLoopHeader->begin();
       PHINode* PN = dyn_cast<PHINode>(I); ++I) {
    if (PN != innerIndexVar) {
      int index = PN->getBasicBlockIndex(innerLoopLatch);
      if (index != -1)
	ValueMap[PN] = PN->getIncomingValue(index);
    }
  }
  // 
  for (BasicBlock::iterator I = outerLoopHeader->begin();
       PHINode* PN = dyn_cast<PHINode>(I); ++I) {
    if (PN != outerIndexVar) {
      int index = PN->getBasicBlockIndex(innerLoopHeader);
      if (index != -1)
	ValueMap[PN] = PN->getIncomingValue(index);
    }
  }
  // Update SSA to get the values from the loop body.
  SSAUpdater SSA;
  for (BasicBlock::iterator I = outerLoopHeader->begin(); 
       PHINode *PN = dyn_cast<PHINode>(I); ) {
    ++I;
    if (PN == outerIndexVar)
      continue;
    Value* outerLoopHeaderValue = PN;
    Value* innerLoopHeaderValue = ValueMap[PN];
    Value* innerLoopLatchPredValue = ValueMap[innerLoopHeaderValue];
    SSA.Initialize(outerLoopHeaderValue);
      SSA.AddAvailableValue(outerLoopHeader, outerLoopHeaderValue);
    if (innerLoopHeaderValue)
      SSA.AddAvailableValue(innerLoopHeader, innerLoopHeaderValue);
    if (innerLoopLatchPredValue)
      SSA.AddAvailableValue(innerLoopLatchPred, innerLoopLatchPredValue);
    for (Value::use_iterator UI = innerLoopHeaderValue->use_begin(),
	   UE = innerLoopHeaderValue->use_end(); UI != UE; ) {
      Use &U = UI.getUse();
      ++UI;
      SSA.RewriteUse(U);
    } 
  }
}

