<commit_before>//===- BottomUpClosure.cpp - Compute bottom-up interprocedural closure ----===//
//
// This file implements the BUDataStructures class, which represents the
// Bottom-Up Interprocedural closure of the data structure graph over the
// program.  This is useful for applications like pool allocation, but **not**
// applications like alias analysis.
//
//===----------------------------------------------------------------------===//

#include "llvm/Analysis/DataStructure.h"
#include "llvm/Analysis/DSGraph.h"
#include "llvm/Module.h"
#include "Support/Statistic.h"
using std::map;

static RegisterAnalysis<BUDataStructures>
X("budatastructure", "Bottom-up Data Structure Analysis Closure");

using namespace DS;

// run - Calculate the bottom up data structure graphs for each function in the
// program.
//
bool BUDataStructures::run(Module &M) {
  GlobalsGraph = new DSGraph();

  // Simply calculate the graphs for each function...
  for (Module::iterator I = M.begin(), E = M.end(); I != E; ++I)
    if (!I->isExternal())
      calculateGraph(*I, 0);
  return false;
}

// releaseMemory - If the pass pipeline is done with this pass, we can release
// our memory... here...
//
void BUDataStructures::releaseMemory() {
  for (map<const Function*, DSGraph*>::iterator I = DSInfo.begin(),
         E = DSInfo.end(); I != E; ++I)
    delete I->second;

  // Empty map so next time memory is released, data structures are not
  // re-deleted.
  DSInfo.clear();
  delete GlobalsGraph;
  GlobalsGraph = 0;
}

DSGraph &BUDataStructures::calculateGraph(Function &F, unsigned Indent) {
  // Make sure this graph has not already been calculated, or that we don't get
  // into an infinite loop with mutually recursive functions.
  //
  DSGraph *&Graph = DSInfo[&F];
  if (Graph) return *Graph;

  // Copy the local version into DSInfo...
  Graph = new DSGraph(getAnalysis<LocalDataStructures>().getDSGraph(F));
  Graph->setGlobalsGraph(GlobalsGraph);
  Graph->setPrintAuxCalls();

  // Start resolving calls...
  std::vector<DSCallSite> &FCs = Graph->getAuxFunctionCalls();

  // Start with a copy of the original call sites...
  FCs = Graph->getFunctionCalls();

  DEBUG(std::cerr << std::string(Indent*4, ' ')
                  << "[BU] Calculating graph for: " << F.getName() << "\n");

  bool Inlined;
  do {
    Inlined = false;

    for (unsigned i = 0; i != FCs.size(); ++i) {
      // Copy the call, because inlining graphs may invalidate the FCs vector.
      DSCallSite Call = FCs[i];

      // If the function list is complete...
      if ((Call.getCallee().getNode()->NodeType & DSNode::Incomplete)==0) {
        // Start inlining all of the functions we can... some may not be
        // inlinable if they are external...
        //
        std::vector<GlobalValue*> Callees =
          Call.getCallee().getNode()->getGlobals();

        unsigned OldNumCalls = FCs.size();

        // Loop over the functions, inlining whatever we can...
        for (unsigned c = 0; c != Callees.size(); ++c) {
          // Must be a function type, so this cast MUST succeed.
          Function &FI = cast<Function>(*Callees[c]);

          if (&FI == &F) {
            // Self recursion... simply link up the formal arguments with the
            // actual arguments...
            DEBUG(std::cerr << std::string(Indent*4, ' ')
                  << "  [BU] Self Inlining: " << F.getName() << "\n");

            // Handle self recursion by resolving the arguments and return value
            Graph->mergeInGraph(Call, *Graph, DSGraph::StripAllocaBit);

            // Erase the entry in the callees vector
            Callees.erase(Callees.begin()+c--);

          } else if (!FI.isExternal()) {
            DEBUG(std::cerr << std::string(Indent*4, ' ')
                  << "  [BU] In " << F.getName() << " inlining: "
                  << FI.getName() << "\n");

            // Get the data structure graph for the called function, closing it
            // if possible (which is only impossible in the case of mutual
            // recursion...
            //
            DSGraph &GI = calculateGraph(FI, Indent+1);  // Graph to inline

            DEBUG(std::cerr << std::string(Indent*4, ' ')
                  << "  [BU] Got graph for " << FI.getName()
                  << " in: " << F.getName() << "[" << GI.getGraphSize() << "+"
                  << GI.getAuxFunctionCalls().size() << "]\n");

            // Handle self recursion by resolving the arguments and return value
            Graph->mergeInGraph(Call, GI, DSGraph::StripAllocaBit |
                                DSGraph::DontCloneCallNodes);

            // Erase the entry in the Callees vector
            Callees.erase(Callees.begin()+c--);

          } else if (FI.getName() == "printf" || FI.getName() == "sscanf" ||
                     FI.getName() == "fprintf" || FI.getName() == "open" ||
                     FI.getName() == "sprintf" || FI.getName() == "fputs") {
            // FIXME: These special cases (eg printf) should go away when we can
            // define functions that take a variable number of arguments.

            // FIXME: at the very least, this should update mod/ref info
            // Erase the entry in the globals vector
            Callees.erase(Callees.begin()+c--);
          }
        }

        if (Callees.empty()) {         // Inlined all of the function calls?
          // Erase the call if it is resolvable...
          FCs.erase(FCs.begin()+i--);  // Don't skip a the next call...
          Inlined = true;
        } else if (Callees.size() !=
                   Call.getCallee().getNode()->getGlobals().size()) {
          // Was able to inline SOME, but not all of the functions.  Construct a
          // new global node here.
          //
          assert(0 && "Unimpl!");
          Inlined = true;
        }

        // If we just inlined a function that had call nodes, chances are that
        // the call nodes are redundant with ones we already have.  Eliminate
        // those call nodes now.
        //
        if (FCs.size() > OldNumCalls)
          Graph->removeTriviallyDeadNodes();
      }

      if (FCs.size() > 200) {
        std::cerr << "Aborted inlining fn: '" << F.getName() << "'!"
                  << std::endl;
        Graph->maskIncompleteMarkers();
        Graph->markIncompleteNodes();
        Graph->removeDeadNodes();
        Graph->writeGraphToFile(std::cerr, "crap."+F.getName());
        exit(1);
        return *Graph;
      }

    }

    // Recompute the Incomplete markers.  If there are any function calls left
    // now that are complete, we must loop!
    if (Inlined) {
      Graph->maskIncompleteMarkers();
      Graph->markIncompleteNodes();
      Graph->removeDeadNodes();
    }

  } while (Inlined && !FCs.empty());

  DEBUG(std::cerr << std::string(Indent*4, ' ')
        << "[BU] Done inlining: " << F.getName() << " ["
        << Graph->getGraphSize() << "+" << Graph->getAuxFunctionCalls().size()
        << "]\n");

  //Graph->writeGraphToFile(std::cerr, "bu_" + F.getName());

  return *Graph;
}
<commit_msg>Fix infinite loop in the BU algorithm.  Unfortunately this dies a serious death when handling moderately sized SCC's, but what can you do<commit_after>//===- BottomUpClosure.cpp - Compute bottom-up interprocedural closure ----===//
//
// This file implements the BUDataStructures class, which represents the
// Bottom-Up Interprocedural closure of the data structure graph over the
// program.  This is useful for applications like pool allocation, but **not**
// applications like alias analysis.
//
//===----------------------------------------------------------------------===//

#include "llvm/Analysis/DataStructure.h"
#include "llvm/Analysis/DSGraph.h"
#include "llvm/Module.h"
#include "Support/Statistic.h"
using std::map;

static RegisterAnalysis<BUDataStructures>
X("budatastructure", "Bottom-up Data Structure Analysis Closure");

using namespace DS;

// run - Calculate the bottom up data structure graphs for each function in the
// program.
//
bool BUDataStructures::run(Module &M) {
  GlobalsGraph = new DSGraph();

  // Simply calculate the graphs for each function...
  for (Module::iterator I = M.begin(), E = M.end(); I != E; ++I)
    if (!I->isExternal())
      calculateGraph(*I, 0);
  return false;
}

// releaseMemory - If the pass pipeline is done with this pass, we can release
// our memory... here...
//
void BUDataStructures::releaseMemory() {
  for (map<const Function*, DSGraph*>::iterator I = DSInfo.begin(),
         E = DSInfo.end(); I != E; ++I)
    delete I->second;

  // Empty map so next time memory is released, data structures are not
  // re-deleted.
  DSInfo.clear();
  delete GlobalsGraph;
  GlobalsGraph = 0;
}

// Return true if a graph was inlined
// Can not modify the part of the AuxCallList < FirstResolvableCall.
//
bool BUDataStructures::ResolveFunctionCalls(DSGraph &G,
                                            unsigned &FirstResolvableCall,
                                   std::map<Function*, DSCallSite> &InProcess,
                                            unsigned Indent) {
  std::vector<DSCallSite> &FCs = G.getAuxFunctionCalls();
  bool Changed = false;

  // Loop while there are call sites that we can resolve!
  while (FirstResolvableCall != FCs.size()) {
    DSCallSite Call = FCs[FirstResolvableCall];

    // If the function list is incomplete...
    if (Call.getCallee().getNode()->NodeType & DSNode::Incomplete) {
      // If incomplete, we cannot resolve it, so leave it at the beginning of
      // the call list with the other unresolvable calls...
      ++FirstResolvableCall;
    } else {
      // Start inlining all of the functions we can... some may not be
      // inlinable if they are external...
      //
      const std::vector<GlobalValue*> &Callees =
        Call.getCallee().getNode()->getGlobals();

      bool hasExternalTarget = false;

      // Loop over the functions, inlining whatever we can...
      for (unsigned c = 0, e = Callees.size(); c != e; ++c) {
        // Must be a function type, so this cast should succeed unless something
        // really wierd is happening.
        Function &FI = cast<Function>(*Callees[c]);

        if (FI.getName() == "printf" || FI.getName() == "sscanf" ||
            FI.getName() == "fprintf" || FI.getName() == "open" ||
            FI.getName() == "sprintf" || FI.getName() == "fputs") {
          // Ignore
        } else if (FI.isExternal()) {
          // If the function is external, then we cannot resolve this call site!
          hasExternalTarget = true;
          break;
        } else {
          std::map<Function*, DSCallSite>::iterator I =
            InProcess.lower_bound(&FI);

          if (I != InProcess.end() && I->first == &FI) {  // Recursion detected?
            // Merge two call sites to eliminate recursion...
            Call.mergeWith(I->second);

            DEBUG(std::cerr << std::string(Indent*2, ' ')
                  << "* Recursion detected for function " << FI.getName()<<"\n");
          } else {
            DEBUG(std::cerr << std::string(Indent*2, ' ')
                  << "Inlining: " << FI.getName() << "\n");

            // Get the data structure graph for the called function, closing it
            // if possible...
            //
            DSGraph &GI = calculateGraph(FI, Indent+1);  // Graph to inline

            DEBUG(std::cerr << std::string(Indent*2, ' ')
                  << "Got graph for: " << FI.getName() << "["
                  << GI.getGraphSize() << "+"
                  << GI.getAuxFunctionCalls().size() << "] "
                  << " in: " << G.getFunction().getName() << "["
                  << G.getGraphSize() << "+"
                  << G.getAuxFunctionCalls().size() << "]\n");

            // Keep track of how many call sites are added by the inlining...
            unsigned NumCalls = FCs.size();

            // Resolve the arguments and return value
            G.mergeInGraph(Call, GI, DSGraph::StripAllocaBit |
                           DSGraph::DontCloneCallNodes);

            // Added a call site?
            if (FCs.size() != NumCalls) {
              // Otherwise we need to inline the graph.  Temporarily add the
              // current function to the InProcess map to be able to handle
              // recursion successfully.
              //
              I = InProcess.insert(I, std::make_pair(&FI, Call));

              // ResolveFunctionCalls - Resolve the function calls that just got
              // inlined...
              //
              Changed |= ResolveFunctionCalls(G, NumCalls, InProcess, Indent+1);

              // Now that we are done processing the inlined graph, remove our
              // cycle detector record...
              //
              //InProcess.erase(I);
            }
          }
        }
      }

      if (hasExternalTarget) {
        // If we cannot resolve this call site...
        ++FirstResolvableCall;
      } else {
        Changed = true;
        FCs.erase(FCs.begin()+FirstResolvableCall);
      }
    }
  }

  return Changed;
}

DSGraph &BUDataStructures::calculateGraph(Function &F, unsigned Indent) {
  // Make sure this graph has not already been calculated, or that we don't get
  // into an infinite loop with mutually recursive functions.
  //
  DSGraph *&GraphPtr = DSInfo[&F];
  if (GraphPtr) return *GraphPtr;

  // Copy the local version into DSInfo...
  GraphPtr = new DSGraph(getAnalysis<LocalDataStructures>().getDSGraph(F));
  DSGraph &Graph = *GraphPtr;

  Graph.setGlobalsGraph(GlobalsGraph);
  Graph.setPrintAuxCalls();

  // Start resolving calls...
  std::vector<DSCallSite> &FCs = Graph.getAuxFunctionCalls();

  // Start with a copy of the original call sites...
  FCs = Graph.getFunctionCalls();

  DEBUG(std::cerr << std::string(Indent*2, ' ')
                  << "[BU] Calculating graph for: " << F.getName() << "\n");

  bool Changed;
  while (1) {
    unsigned FirstResolvableCall = 0;
    std::map<Function *, DSCallSite> InProcess;

    // Insert a call site for self to handle self recursion...
    std::vector<DSNodeHandle> Args;
    Args.reserve(F.asize());
    for (Function::aiterator I = F.abegin(), E = F.aend(); I != E; ++I)
      if (isPointerType(I->getType()))
        Args.push_back(Graph.getNodeForValue(I));

    InProcess.insert(std::make_pair(&F,
           DSCallSite(*(CallInst*)0, Graph.getRetNode(),(DSNode*)0,Args)));

    Changed = ResolveFunctionCalls(Graph, FirstResolvableCall, InProcess,
                                   Indent);

    if (Changed) {
      Graph.maskIncompleteMarkers();
      Graph.markIncompleteNodes();
      Graph.removeDeadNodes();
      break;
    } else {
      break;
    }
  }

#if 0
  bool Inlined;
  do {
    Inlined = false;

    for (unsigned i = 0; i != FCs.size(); ++i) {
      // Copy the call, because inlining graphs may invalidate the FCs vector.
      DSCallSite Call = FCs[i];

      // If the function list is complete...
      if ((Call.getCallee().getNode()->NodeType & DSNode::Incomplete)==0) {
        // Start inlining all of the functions we can... some may not be
        // inlinable if they are external...
        //
        std::vector<GlobalValue*> Callees =
          Call.getCallee().getNode()->getGlobals();

        unsigned OldNumCalls = FCs.size();

        // Loop over the functions, inlining whatever we can...
        for (unsigned c = 0; c != Callees.size(); ++c) {
          // Must be a function type, so this cast MUST succeed.
          Function &FI = cast<Function>(*Callees[c]);

          if (&FI == &F) {
            // Self recursion... simply link up the formal arguments with the
            // actual arguments...
            DEBUG(std::cerr << std::string(Indent*2, ' ')
                  << "[BU] Self Inlining: " << F.getName() << "\n");

            // Handle self recursion by resolving the arguments and return value
            Graph.mergeInGraph(Call, Graph, DSGraph::StripAllocaBit);

            // Erase the entry in the callees vector
            Callees.erase(Callees.begin()+c--);

          } else if (!FI.isExternal()) {
            DEBUG(std::cerr << std::string(Indent*2, ' ')
                  << "[BU] In " << F.getName() << " inlining: "
                  << FI.getName() << "\n");

            // Get the data structure graph for the called function, closing it
            // if possible (which is only impossible in the case of mutual
            // recursion...
            //
            DSGraph &GI = calculateGraph(FI, Indent+1);  // Graph to inline

            DEBUG(std::cerr << std::string(Indent*2, ' ')
                  << "[BU] Got graph for " << FI.getName()
                  << " in: " << F.getName() << "[" << GI.getGraphSize() << "+"
                  << GI.getAuxFunctionCalls().size() << "]\n");

            // Handle self recursion by resolving the arguments and return value
            Graph.mergeInGraph(Call, GI, DSGraph::StripAllocaBit |
                                DSGraph::DontCloneCallNodes);

            // Erase the entry in the Callees vector
            Callees.erase(Callees.begin()+c--);

          } else if (FI.getName() == "printf" || FI.getName() == "sscanf" ||
                     FI.getName() == "fprintf" || FI.getName() == "open" ||
                     FI.getName() == "sprintf" || FI.getName() == "fputs") {
            // FIXME: These special cases (eg printf) should go away when we can
            // define functions that take a variable number of arguments.

            // FIXME: at the very least, this should update mod/ref info
            // Erase the entry in the globals vector
            Callees.erase(Callees.begin()+c--);
          }
        }

        if (Callees.empty()) {         // Inlined all of the function calls?
          // Erase the call if it is resolvable...
          FCs.erase(FCs.begin()+i--);  // Don't skip a the next call...
          Inlined = true;
        } else if (Callees.size() !=
                   Call.getCallee().getNode()->getGlobals().size()) {
          // Was able to inline SOME, but not all of the functions.  Construct a
          // new global node here.
          //
          assert(0 && "Unimpl!");
          Inlined = true;
        }

#if 0
        // If we just inlined a function that had call nodes, chances are that
        // the call nodes are redundant with ones we already have.  Eliminate
        // those call nodes now.
        //
        if (FCs.size() >= OldNumCalls)
          Graph.removeTriviallyDeadNodes();
#endif
      }

      if (FCs.size() > 200) {
        std::cerr << "Aborted inlining fn: '" << F.getName() << "'!"
                  << std::endl;
        Graph.maskIncompleteMarkers();
        Graph.markIncompleteNodes();
        Graph.removeDeadNodes();
        Graph.writeGraphToFile(std::cerr, "crap."+F.getName());
        exit(1);
        return Graph;
      }

    }

    // Recompute the Incomplete markers.  If there are any function calls left
    // now that are complete, we must loop!
    if (Inlined) {
      Graph.maskIncompleteMarkers();
      Graph.markIncompleteNodes();
      Graph.removeDeadNodes();
    }

  } while (Inlined && !FCs.empty());
#endif

  DEBUG(std::cerr << std::string(Indent*2, ' ')
        << "[BU] Done inlining: " << F.getName() << " ["
        << Graph.getGraphSize() << "+" << Graph.getAuxFunctionCalls().size()
        << "]\n");

  Graph.writeGraphToFile(std::cerr, "bu_" + F.getName());

  return Graph;
}
<|endoftext|>


<commit_before>//===--- ClangDerivedConformances.cpp -------------------------------------===//
//
// This source file is part of the Swift.org open source project
//
// Copyright (c) 2014 - 2022 Apple Inc. and the Swift project authors
// Licensed under Apache License v2.0 with Runtime Library Exception
//
// See https://swift.org/LICENSE.txt for license information
// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
//
//===----------------------------------------------------------------------===//

#include "ClangDerivedConformances.h"
#include "swift/AST/NameLookup.h"
#include "swift/AST/ParameterList.h"

using namespace swift;

static clang::TypeDecl *
getIteratorCategoryDecl(const clang::CXXRecordDecl *clangDecl) {
  clang::IdentifierInfo *iteratorCategoryDeclName =
      &clangDecl->getASTContext().Idents.get("iterator_category");
  auto iteratorCategories = clangDecl->lookup(iteratorCategoryDeclName);
  if (!iteratorCategories.isSingleResult())
    return nullptr;
  auto iteratorCategory = iteratorCategories.front();

  return dyn_cast_or_null<clang::TypeDecl>(iteratorCategory);
}

static ValueDecl *getEqualEqualOperator(NominalTypeDecl *decl) {
  auto id = decl->getASTContext().Id_EqualsOperator;

  auto isValid = [&](ValueDecl *equalEqualOp) -> bool {
    auto equalEqual = dyn_cast<FuncDecl>(equalEqualOp);
    if (!equalEqual || !equalEqual->hasParameterList())
      return false;
    auto params = equalEqual->getParameters();
    if (params->size() != 2)
      return false;
    auto lhs = params->get(0);
    auto rhs = params->get(1);
    if (lhs->isInOut() || rhs->isInOut())
      return false;
    auto lhsTy = lhs->getType();
    auto rhsTy = rhs->getType();
    if (!lhsTy || !rhsTy)
      return false;
    auto lhsNominal = lhsTy->getAnyNominal();
    auto rhsNominal = rhsTy->getAnyNominal();
    if (lhsNominal != rhsNominal || lhsNominal != decl)
      return false;
    return true;
  };

  // First look for `func ==` declared as a member.
  auto memberResults = decl->lookupDirect(id);
  for (const auto &member : memberResults) {
    if (isValid(member))
      return member;
  }

  // If no member `func ==` was found, look for out-of-class definitions in the
  // same module.
  auto module = decl->getModuleContext();
  llvm::SmallVector<ValueDecl *> nonMemberResults;
  module->lookupValue(id, NLKind::UnqualifiedLookup, nonMemberResults);
  for (const auto &nonMember : nonMemberResults) {
    if (isValid(nonMember))
      return nonMember;
  }

  return nullptr;
}

bool swift::isIterator(const clang::CXXRecordDecl *clangDecl) {
  return getIteratorCategoryDecl(clangDecl);
}

void swift::conformToCxxIteratorIfNeeded(
    ClangImporter::Implementation &impl, NominalTypeDecl *decl,
    const clang::CXXRecordDecl *clangDecl) {
  assert(decl);
  assert(clangDecl);
  ASTContext &ctx = decl->getASTContext();

  // We consider a type to be an input iterator if it defines an
  // `iterator_category` that inherits from `std::input_iterator_tag`, e.g.
  // `using iterator_category = std::input_iterator_tag`.
  auto iteratorCategory = getIteratorCategoryDecl(clangDecl);
  if (!iteratorCategory)
    return;

  // If `iterator_category` is a typedef or a using-decl, retrieve the
  // underlying struct decl.
  clang::CXXRecordDecl *underlyingCategoryDecl = nullptr;
  if (auto typedefDecl = dyn_cast<clang::TypedefNameDecl>(iteratorCategory)) {
    auto type = typedefDecl->getUnderlyingType();
    underlyingCategoryDecl = type->getAsCXXRecordDecl();
  } else {
    underlyingCategoryDecl = dyn_cast<clang::CXXRecordDecl>(iteratorCategory);
  }
  if (underlyingCategoryDecl) {
    underlyingCategoryDecl = underlyingCategoryDecl->getDefinition();
  }

  if (!underlyingCategoryDecl)
    return;

  auto isInputIteratorDecl = [&](const clang::CXXRecordDecl *base) {
    return base->isInStdNamespace() && base->getIdentifier() &&
           base->getName() == "input_iterator_tag";
  };

  // Traverse all transitive bases of `underlyingDecl` to check if
  // it inherits from `std::input_iterator_tag`.
  bool isInputIterator = isInputIteratorDecl(underlyingCategoryDecl);
  underlyingCategoryDecl->forallBases([&](const clang::CXXRecordDecl *base) {
    if (isInputIteratorDecl(base)) {
      isInputIterator = true;
      return false;
    }
    return true;
  });

  if (!isInputIterator)
    return;

  // Check if present: `var pointee: Pointee { get }`
  auto pointeeId = ctx.getIdentifier("pointee");
  auto pointees = decl->lookupDirect(pointeeId);
  if (pointees.size() != 1)
    return;
  auto pointee = dyn_cast<VarDecl>(pointees.front());
  if (!pointee || pointee->isGetterMutating())
    return;

  // Check if present: `func successor() -> Self`
  auto successorId = ctx.getIdentifier("successor");
  auto successors = decl->lookupDirect(successorId);
  if (successors.size() != 1)
    return;
  auto successor = dyn_cast<FuncDecl>(successors.front());
  if (!successor || successor->isMutating())
    return;
  auto successorTy = successor->getResultInterfaceType();
  if (!successorTy || successorTy->getAnyNominal() != decl)
    return;

  // Check if present: `func ==`
  auto equalEqual = getEqualEqualOperator(decl);
  if (!equalEqual)
    return;

  impl.addSynthesizedTypealias(decl, ctx.getIdentifier("Pointee"),
                               pointee->getType());
  impl.addSynthesizedProtocolAttrs(decl,
                                   {KnownProtocolKind::UnsafeCxxInputIterator});
}
<commit_msg>[cxx-interop][NFC] Prettify stack trace for CxxIterator conformance synthesis<commit_after>//===--- ClangDerivedConformances.cpp -------------------------------------===//
//
// This source file is part of the Swift.org open source project
//
// Copyright (c) 2014 - 2022 Apple Inc. and the Swift project authors
// Licensed under Apache License v2.0 with Runtime Library Exception
//
// See https://swift.org/LICENSE.txt for license information
// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
//
//===----------------------------------------------------------------------===//

#include "ClangDerivedConformances.h"
#include "swift/AST/NameLookup.h"
#include "swift/AST/ParameterList.h"
#include "swift/AST/PrettyStackTrace.h"

using namespace swift;

static clang::TypeDecl *
getIteratorCategoryDecl(const clang::CXXRecordDecl *clangDecl) {
  clang::IdentifierInfo *iteratorCategoryDeclName =
      &clangDecl->getASTContext().Idents.get("iterator_category");
  auto iteratorCategories = clangDecl->lookup(iteratorCategoryDeclName);
  if (!iteratorCategories.isSingleResult())
    return nullptr;
  auto iteratorCategory = iteratorCategories.front();

  return dyn_cast_or_null<clang::TypeDecl>(iteratorCategory);
}

static ValueDecl *getEqualEqualOperator(NominalTypeDecl *decl) {
  auto id = decl->getASTContext().Id_EqualsOperator;

  auto isValid = [&](ValueDecl *equalEqualOp) -> bool {
    auto equalEqual = dyn_cast<FuncDecl>(equalEqualOp);
    if (!equalEqual || !equalEqual->hasParameterList())
      return false;
    auto params = equalEqual->getParameters();
    if (params->size() != 2)
      return false;
    auto lhs = params->get(0);
    auto rhs = params->get(1);
    if (lhs->isInOut() || rhs->isInOut())
      return false;
    auto lhsTy = lhs->getType();
    auto rhsTy = rhs->getType();
    if (!lhsTy || !rhsTy)
      return false;
    auto lhsNominal = lhsTy->getAnyNominal();
    auto rhsNominal = rhsTy->getAnyNominal();
    if (lhsNominal != rhsNominal || lhsNominal != decl)
      return false;
    return true;
  };

  // First look for `func ==` declared as a member.
  auto memberResults = decl->lookupDirect(id);
  for (const auto &member : memberResults) {
    if (isValid(member))
      return member;
  }

  // If no member `func ==` was found, look for out-of-class definitions in the
  // same module.
  auto module = decl->getModuleContext();
  llvm::SmallVector<ValueDecl *> nonMemberResults;
  module->lookupValue(id, NLKind::UnqualifiedLookup, nonMemberResults);
  for (const auto &nonMember : nonMemberResults) {
    if (isValid(nonMember))
      return nonMember;
  }

  return nullptr;
}

bool swift::isIterator(const clang::CXXRecordDecl *clangDecl) {
  return getIteratorCategoryDecl(clangDecl);
}

void swift::conformToCxxIteratorIfNeeded(
    ClangImporter::Implementation &impl, NominalTypeDecl *decl,
    const clang::CXXRecordDecl *clangDecl) {
  PrettyStackTraceDecl trace("conforming to UnsafeCxxInputIterator", decl);

  assert(decl);
  assert(clangDecl);
  ASTContext &ctx = decl->getASTContext();

  // We consider a type to be an input iterator if it defines an
  // `iterator_category` that inherits from `std::input_iterator_tag`, e.g.
  // `using iterator_category = std::input_iterator_tag`.
  auto iteratorCategory = getIteratorCategoryDecl(clangDecl);
  if (!iteratorCategory)
    return;

  // If `iterator_category` is a typedef or a using-decl, retrieve the
  // underlying struct decl.
  clang::CXXRecordDecl *underlyingCategoryDecl = nullptr;
  if (auto typedefDecl = dyn_cast<clang::TypedefNameDecl>(iteratorCategory)) {
    auto type = typedefDecl->getUnderlyingType();
    underlyingCategoryDecl = type->getAsCXXRecordDecl();
  } else {
    underlyingCategoryDecl = dyn_cast<clang::CXXRecordDecl>(iteratorCategory);
  }
  if (underlyingCategoryDecl) {
    underlyingCategoryDecl = underlyingCategoryDecl->getDefinition();
  }

  if (!underlyingCategoryDecl)
    return;

  auto isInputIteratorDecl = [&](const clang::CXXRecordDecl *base) {
    return base->isInStdNamespace() && base->getIdentifier() &&
           base->getName() == "input_iterator_tag";
  };

  // Traverse all transitive bases of `underlyingDecl` to check if
  // it inherits from `std::input_iterator_tag`.
  bool isInputIterator = isInputIteratorDecl(underlyingCategoryDecl);
  underlyingCategoryDecl->forallBases([&](const clang::CXXRecordDecl *base) {
    if (isInputIteratorDecl(base)) {
      isInputIterator = true;
      return false;
    }
    return true;
  });

  if (!isInputIterator)
    return;

  // Check if present: `var pointee: Pointee { get }`
  auto pointeeId = ctx.getIdentifier("pointee");
  auto pointees = decl->lookupDirect(pointeeId);
  if (pointees.size() != 1)
    return;
  auto pointee = dyn_cast<VarDecl>(pointees.front());
  if (!pointee || pointee->isGetterMutating())
    return;

  // Check if present: `func successor() -> Self`
  auto successorId = ctx.getIdentifier("successor");
  auto successors = decl->lookupDirect(successorId);
  if (successors.size() != 1)
    return;
  auto successor = dyn_cast<FuncDecl>(successors.front());
  if (!successor || successor->isMutating())
    return;
  auto successorTy = successor->getResultInterfaceType();
  if (!successorTy || successorTy->getAnyNominal() != decl)
    return;

  // Check if present: `func ==`
  auto equalEqual = getEqualEqualOperator(decl);
  if (!equalEqual)
    return;

  impl.addSynthesizedTypealias(decl, ctx.getIdentifier("Pointee"),
                               pointee->getType());
  impl.addSynthesizedProtocolAttrs(decl,
                                   {KnownProtocolKind::UnsafeCxxInputIterator});
}
<|endoftext|>


<commit_before>/*=========================================================================
 *
 *  Copyright RTK Consortium
 *
 *  Licensed under the Apache License, Version 2.0 (the "License");
 *  you may not use this file except in compliance with the License.
 *  You may obtain a copy of the License at
 *
 *         http://www.apache.org/licenses/LICENSE-2.0.txt
 *
 *  Unless required by applicable law or agreed to in writing, software
 *  distributed under the License is distributed on an "AS IS" BASIS,
 *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *  See the License for the specific language governing permissions and
 *  limitations under the License.
 *
 *=========================================================================*/

#include "rtkfdk_ggo.h"
#include "rtkGgoFunctions.h"
#include "rtkConfiguration.h"

#include "rtkThreeDCircularProjectionGeometryXMLFile.h"
#include "rtkProjectionsReader.h"
#include "rtkDisplacedDetectorImageFilter.h"
#include "rtkParkerShortScanImageFilter.h"
#include "rtkFDKConeBeamReconstructionFilter.h"
#if CUDA_FOUND
# include "rtkCudaFDKConeBeamReconstructionFilter.h"
#endif
#if OPENCL_FOUND
# include "rtkOpenCLFDKConeBeamReconstructionFilter.h"
#endif
#include "rtkFDKWarpBackProjectionImageFilter.h"
#include "rtkCyclicDeformationImageFilter.h"

#include <itkRegularExpressionSeriesFileNames.h>
#include <itkStreamingImageFilter.h>
#include <itkImageFileWriter.h>

int main(int argc, char * argv[])
{
  GGO(rtkfdk, args_info);

  typedef float OutputPixelType;
  const unsigned int Dimension = 3;

#if CUDA_FOUND
  typedef itk::CudaImage< OutputPixelType, Dimension > OutputImageType;
#else
  typedef itk::Image< OutputPixelType, Dimension > OutputImageType;
#endif

  // Generate file names
  itk::RegularExpressionSeriesFileNames::Pointer names = itk::RegularExpressionSeriesFileNames::New();
  names->SetDirectory(args_info.path_arg);
  names->SetNumericSort(false);
  names->SetRegularExpression(args_info.regexp_arg);
  names->SetSubMatch(0);

  if(args_info.verbose_flag)
    std::cout << "Regular expression matches "
              << names->GetFileNames().size()
              << " file(s)..."
              << std::endl;

  // Projections reader
  typedef rtk::ProjectionsReader< OutputImageType > ReaderType;
  ReaderType::Pointer reader = ReaderType::New();
  reader->SetFileNames( names->GetFileNames() );
  TRY_AND_EXIT_ON_ITK_EXCEPTION( reader->GenerateOutputInformation() );

  itk::TimeProbe readerProbe;
  if(!args_info.lowmem_flag)
    {
    if(args_info.verbose_flag)
      std::cout << "Reading... " << std::flush;
    readerProbe.Start();
    TRY_AND_EXIT_ON_ITK_EXCEPTION( reader->Update() )
    readerProbe.Stop();
    if(args_info.verbose_flag)
      std::cout << "It took " << readerProbe.GetMean() << ' ' << readerProbe.GetUnit() << std::endl;
    }

  // Geometry
  if(args_info.verbose_flag)
    std::cout << "Reading geometry information from "
              << args_info.geometry_arg
              << "..."
              << std::endl;
  rtk::ThreeDCircularProjectionGeometryXMLFileReader::Pointer geometryReader;
  geometryReader = rtk::ThreeDCircularProjectionGeometryXMLFileReader::New();
  geometryReader->SetFilename(args_info.geometry_arg);
  TRY_AND_EXIT_ON_ITK_EXCEPTION( geometryReader->GenerateOutputInformation() )

  // Displaced detector weighting
  typedef rtk::DisplacedDetectorImageFilter< OutputImageType > DDFType;
  DDFType::Pointer ddf = DDFType::New();
  ddf->SetInput( reader->GetOutput() );
  ddf->SetGeometry( geometryReader->GetOutputObject() );

  // Short scan image filter
  typedef rtk::ParkerShortScanImageFilter< OutputImageType > PSSFType;
  PSSFType::Pointer pssf = PSSFType::New();
  pssf->SetInput( ddf->GetOutput() );
  pssf->SetGeometry( geometryReader->GetOutputObject() );
  pssf->InPlaceOff();

  // Create reconstructed image
  typedef rtk::ConstantImageSource< OutputImageType > ConstantImageSourceType;
  ConstantImageSourceType::Pointer constantImageSource = ConstantImageSourceType::New();
  rtk::SetConstantImageSourceFromGgo<ConstantImageSourceType, args_info_rtkfdk>(constantImageSource, args_info);

  // Motion-compensated objects for the compensation of a cyclic deformation.
  // Although these will only be used if the command line options for motion
  // compensation are set, we still create the object before hand to avoid auto
  // destruction.
  typedef itk::Vector<float,3> DVFPixelType;
  typedef itk::Image< DVFPixelType, 3 > DVFImageType;
  typedef rtk::CyclicDeformationImageFilter< DVFImageType > DeformationType;
  typedef itk::ImageFileReader<DeformationType::InputImageType> DVFReaderType;
  DVFReaderType::Pointer dvfReader = DVFReaderType::New();
  DeformationType::Pointer def = DeformationType::New();
  def->SetInput(dvfReader->GetOutput());
  typedef rtk::FDKWarpBackProjectionImageFilter<OutputImageType, OutputImageType, DeformationType> WarpBPType;
  WarpBPType::Pointer bp = WarpBPType::New();
  bp->SetDeformation(def);
  bp->SetGeometry( geometryReader->GetOutputObject() );

  // This macro sets options for fdk filter which I can not see how to do better
  // because TFFTPrecision is not the same, e.g. for CPU and CUDA (SR)
#define SET_FELDKAMP_OPTIONS(f) \
    f->SetInput( 0, constantImageSource->GetOutput() ); \
    f->SetInput( 1, pssf->GetOutput() ); \
    f->SetGeometry( geometryReader->GetOutputObject() ); \
    f->GetRampFilter()->SetTruncationCorrection(args_info.pad_arg); \
    f->GetRampFilter()->SetHannCutFrequency(args_info.hann_arg); \
    f->GetRampFilter()->SetHannCutFrequencyY(args_info.hannY_arg);

  // FDK reconstruction filtering
  itk::ImageToImageFilter<OutputImageType, OutputImageType>::Pointer feldkamp;
  typedef rtk::FDKConeBeamReconstructionFilter< OutputImageType > FDKCPUType;
#if CUDA_FOUND
  typedef rtk::CudaFDKConeBeamReconstructionFilter                FDKCUDAType;
#endif
#if OPENCL_FOUND
  typedef rtk::OpenCLFDKConeBeamReconstructionFilter              FDKOPENCLType;
#endif
  if(!strcmp(args_info.hardware_arg, "cpu") )
    {
    feldkamp = FDKCPUType::New();
    SET_FELDKAMP_OPTIONS( dynamic_cast<FDKCPUType*>(feldkamp.GetPointer()) );

    // Motion compensated CBCT settings
    if(args_info.signal_given && args_info.dvf_given)
      {
      dvfReader->SetFileName(args_info.dvf_arg);
      def->SetSignalFilename(args_info.signal_arg);
      dynamic_cast<FDKCPUType*>(feldkamp.GetPointer())->SetBackProjectionFilter( bp.GetPointer() );
      }
    }
  else if(!strcmp(args_info.hardware_arg, "cuda") )
    {
#if CUDA_FOUND
    feldkamp = FDKCUDAType::New();
    SET_FELDKAMP_OPTIONS( static_cast<FDKCUDAType*>(feldkamp.GetPointer()) );
#else
    std::cerr << "The program has not been compiled with cuda option" << std::endl;
    return EXIT_FAILURE;
#endif
    }
  else if(!strcmp(args_info.hardware_arg, "opencl") )
    {
#if OPENCL_FOUND
    feldkamp = FDKOPENCLType::New();
    SET_FELDKAMP_OPTIONS( static_cast<FDKOPENCLType*>(feldkamp.GetPointer()) );
#else
    std::cerr << "The program has not been compiled with opencl option" << std::endl;
    return EXIT_FAILURE;
#endif
    }

  // Streaming depending on streaming capability of writer
  typedef itk::StreamingImageFilter<OutputImageType, OutputImageType> StreamerType;
  StreamerType::Pointer streamerBP = StreamerType::New();
  streamerBP->SetInput( feldkamp->GetOutput() );
  streamerBP->SetNumberOfStreamDivisions( args_info.divisions_arg );

  // Write
  typedef itk::ImageFileWriter<  OutputImageType > WriterType;
  WriterType::Pointer writer = WriterType::New();
  writer->SetFileName( args_info.output_arg );
  writer->SetInput( streamerBP->GetOutput() );

  if(args_info.verbose_flag)
    std::cout << "Reconstructing and writing... " << std::flush;
  itk::TimeProbe writerProbe;

  writerProbe.Start();
  TRY_AND_EXIT_ON_ITK_EXCEPTION( writer->Update() );
  writerProbe.Stop();

  if(args_info.verbose_flag)
    {
    std::cout << "It took " << writerProbe.GetMean() << ' ' << readerProbe.GetUnit() << std::endl;
    if(!strcmp(args_info.hardware_arg, "cpu") )
      static_cast<FDKCPUType* >(feldkamp.GetPointer())->PrintTiming(std::cout);
#if CUDA_FOUND
    else if(!strcmp(args_info.hardware_arg, "cuda") )
      static_cast<FDKCUDAType*>(feldkamp.GetPointer())->PrintTiming(std::cout);
#endif
#if OPENCL_FOUND
    else if(!strcmp(args_info.hardware_arg, "opencl") )
      static_cast<FDKOPENCLType*>(feldkamp.GetPointer())->PrintTiming(std::cout);
#endif
    std::cout << std::endl;
    }

  return EXIT_SUCCESS;
}
<commit_msg>Adapted rtkfdk to new cuda image type<commit_after>/*=========================================================================
 *
 *  Copyright RTK Consortium
 *
 *  Licensed under the Apache License, Version 2.0 (the "License");
 *  you may not use this file except in compliance with the License.
 *  You may obtain a copy of the License at
 *
 *         http://www.apache.org/licenses/LICENSE-2.0.txt
 *
 *  Unless required by applicable law or agreed to in writing, software
 *  distributed under the License is distributed on an "AS IS" BASIS,
 *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *  See the License for the specific language governing permissions and
 *  limitations under the License.
 *
 *=========================================================================*/

#include "rtkfdk_ggo.h"
#include "rtkGgoFunctions.h"
#include "rtkConfiguration.h"

#include "rtkThreeDCircularProjectionGeometryXMLFile.h"
#include "rtkProjectionsReader.h"
#include "rtkDisplacedDetectorImageFilter.h"
#include "rtkParkerShortScanImageFilter.h"
#include "rtkFDKConeBeamReconstructionFilter.h"
#if CUDA_FOUND
# include "rtkCudaFDKConeBeamReconstructionFilter.h"
#endif
#if OPENCL_FOUND
# include "rtkOpenCLFDKConeBeamReconstructionFilter.h"
#endif
#include "rtkFDKWarpBackProjectionImageFilter.h"
#include "rtkCyclicDeformationImageFilter.h"

#include <itkRegularExpressionSeriesFileNames.h>
#include <itkStreamingImageFilter.h>
#include <itkImageFileWriter.h>

int main(int argc, char * argv[])
{
  GGO(rtkfdk, args_info);

  typedef float OutputPixelType;
  const unsigned int Dimension = 3;

  typedef itk::Image< OutputPixelType, Dimension >     CPUOutputImageType;
#if CUDA_FOUND
  typedef itk::CudaImage< OutputPixelType, Dimension > OutputImageType;
#else
  typedef CPUOutputImageType                           OutputImageType;
#endif

  // Generate file names
  itk::RegularExpressionSeriesFileNames::Pointer names = itk::RegularExpressionSeriesFileNames::New();
  names->SetDirectory(args_info.path_arg);
  names->SetNumericSort(false);
  names->SetRegularExpression(args_info.regexp_arg);
  names->SetSubMatch(0);

  if(args_info.verbose_flag)
    std::cout << "Regular expression matches "
              << names->GetFileNames().size()
              << " file(s)..."
              << std::endl;

  // Projections reader
  typedef rtk::ProjectionsReader< OutputImageType > ReaderType;
  ReaderType::Pointer reader = ReaderType::New();
  reader->SetFileNames( names->GetFileNames() );
  TRY_AND_EXIT_ON_ITK_EXCEPTION( reader->GenerateOutputInformation() );

  itk::TimeProbe readerProbe;
  if(!args_info.lowmem_flag)
    {
    if(args_info.verbose_flag)
      std::cout << "Reading... " << std::flush;
    readerProbe.Start();
    TRY_AND_EXIT_ON_ITK_EXCEPTION( reader->Update() )
    readerProbe.Stop();
    if(args_info.verbose_flag)
      std::cout << "It took " << readerProbe.GetMean() << ' ' << readerProbe.GetUnit() << std::endl;
    }

  // Geometry
  if(args_info.verbose_flag)
    std::cout << "Reading geometry information from "
              << args_info.geometry_arg
              << "..."
              << std::endl;
  rtk::ThreeDCircularProjectionGeometryXMLFileReader::Pointer geometryReader;
  geometryReader = rtk::ThreeDCircularProjectionGeometryXMLFileReader::New();
  geometryReader->SetFilename(args_info.geometry_arg);
  TRY_AND_EXIT_ON_ITK_EXCEPTION( geometryReader->GenerateOutputInformation() )

  // Displaced detector weighting
  typedef rtk::DisplacedDetectorImageFilter< OutputImageType > DDFType;
  DDFType::Pointer ddf = DDFType::New();
  ddf->SetInput( reader->GetOutput() );
  ddf->SetGeometry( geometryReader->GetOutputObject() );

  // Short scan image filter
  typedef rtk::ParkerShortScanImageFilter< OutputImageType > PSSFType;
  PSSFType::Pointer pssf = PSSFType::New();
  pssf->SetInput( ddf->GetOutput() );
  pssf->SetGeometry( geometryReader->GetOutputObject() );
  pssf->InPlaceOff();

  // Create reconstructed image
  typedef rtk::ConstantImageSource< OutputImageType > ConstantImageSourceType;
  ConstantImageSourceType::Pointer constantImageSource = ConstantImageSourceType::New();
  rtk::SetConstantImageSourceFromGgo<ConstantImageSourceType, args_info_rtkfdk>(constantImageSource, args_info);

  // Motion-compensated objects for the compensation of a cyclic deformation.
  // Although these will only be used if the command line options for motion
  // compensation are set, we still create the object before hand to avoid auto
  // destruction.
  typedef itk::Vector<float,3> DVFPixelType;
  typedef itk::Image< DVFPixelType, 3 > DVFImageType;
  typedef rtk::CyclicDeformationImageFilter< DVFImageType > DeformationType;
  typedef itk::ImageFileReader<DeformationType::InputImageType> DVFReaderType;
  DVFReaderType::Pointer dvfReader = DVFReaderType::New();
  DeformationType::Pointer def = DeformationType::New();
  def->SetInput(dvfReader->GetOutput());
  typedef rtk::FDKWarpBackProjectionImageFilter<OutputImageType, OutputImageType, DeformationType> WarpBPType;
  WarpBPType::Pointer bp = WarpBPType::New();
  bp->SetDeformation(def);
  bp->SetGeometry( geometryReader->GetOutputObject() );

  // This macro sets options for fdk filter which I can not see how to do better
  // because TFFTPrecision is not the same, e.g. for CPU and CUDA (SR)
#define SET_FELDKAMP_OPTIONS(f) \
    f->SetInput( 0, constantImageSource->GetOutput() ); \
    f->SetInput( 1, pssf->GetOutput() ); \
    f->SetGeometry( geometryReader->GetOutputObject() ); \
    f->GetRampFilter()->SetTruncationCorrection(args_info.pad_arg); \
    f->GetRampFilter()->SetHannCutFrequency(args_info.hann_arg); \
    f->GetRampFilter()->SetHannCutFrequencyY(args_info.hannY_arg);

  // FDK reconstruction filtering
  typedef rtk::FDKConeBeamReconstructionFilter< OutputImageType > FDKCPUType;
  FDKCPUType::Pointer feldkamp = FDKCPUType::New();
#if OPENCL_FOUND
  typedef rtk::OpenCLFDKConeBeamReconstructionFilter FDKOPENCLType;
  FDKOPENCLType::Pointer feldkampOCL = FDKOPENCLType::New();
#endif
#if CUDA_FOUND
  typedef rtk::CudaFDKConeBeamReconstructionFilter FDKCUDAType;
  FDKCUDAType::Pointer feldkampCUDA = FDKCUDAType::New();
#endif
  itk::Image< OutputPixelType, Dimension > *pfeldkamp;
  if(!strcmp(args_info.hardware_arg, "cpu") )
    {
    typedef rtk::FDKConeBeamReconstructionFilter< OutputImageType > FDKCPUType;
    feldkamp = FDKCPUType::New();
    SET_FELDKAMP_OPTIONS( feldkamp );

    // Motion compensated CBCT settings
    if(args_info.signal_given && args_info.dvf_given)
      {
      dvfReader->SetFileName(args_info.dvf_arg);
      def->SetSignalFilename(args_info.signal_arg);
      feldkamp->SetBackProjectionFilter( bp.GetPointer() );
      }
    pfeldkamp = feldkamp->GetOutput();
    }
  else if(!strcmp(args_info.hardware_arg, "cuda") )
    {
#if CUDA_FOUND
    SET_FELDKAMP_OPTIONS( feldkampCUDA );
    pfeldkamp = feldkampCUDA->GetOutput();
#else
    std::cerr << "The program has not been compiled with cuda option" << std::endl;
    return EXIT_FAILURE;
#endif
    }
  else if(!strcmp(args_info.hardware_arg, "opencl") )
    {
#if OPENCL_FOUND
    SET_FELDKAMP_OPTIONS( feldkampOCL );
    pfeldkamp = feldkampOCL->GetOutput();
#else
    std::cerr << "The program has not been compiled with opencl option" << std::endl;
    return EXIT_FAILURE;
#endif
    }

  // Streaming depending on streaming capability of writer
  typedef itk::StreamingImageFilter<CPUOutputImageType, CPUOutputImageType> StreamerType;
  StreamerType::Pointer streamerBP = StreamerType::New();
  streamerBP->SetInput( pfeldkamp );
  streamerBP->SetNumberOfStreamDivisions( args_info.divisions_arg );

  // Write
  typedef itk::ImageFileWriter<CPUOutputImageType> WriterType;
  WriterType::Pointer writer = WriterType::New();
  writer->SetFileName( args_info.output_arg );
  writer->SetInput( streamerBP->GetOutput() );

  if(args_info.verbose_flag)
    std::cout << "Reconstructing and writing... " << std::flush;
  itk::TimeProbe writerProbe;

  writerProbe.Start();
  TRY_AND_EXIT_ON_ITK_EXCEPTION( writer->Update() );
  writerProbe.Stop();

  if(args_info.verbose_flag)
    {
    std::cout << "It took " << writerProbe.GetMean() << ' ' << readerProbe.GetUnit() << std::endl;
    if(!strcmp(args_info.hardware_arg, "cpu") )
      feldkamp->PrintTiming(std::cout);
#if CUDA_FOUND
    else if(!strcmp(args_info.hardware_arg, "cuda") )
      feldkampCUDA->PrintTiming(std::cout);
#endif
#if OPENCL_FOUND
    else if(!strcmp(args_info.hardware_arg, "opencl") )
      feldkampOCL->PrintTiming(std::cout);
#endif
    std::cout << std::endl;
    }

  return EXIT_SUCCESS;
}
<|endoftext|>


<commit_before>////////////////////////////////////////////////////////////////////////////////
/// DISCLAIMER
///
/// Copyright 2016 ArangoDB GmbH, Cologne, Germany
///
/// Licensed under the Apache License, Version 2.0 (the "License");
/// you may not use this file except in compliance with the License.
/// You may obtain a copy of the License at
///
///     http://www.apache.org/licenses/LICENSE-2.0
///
/// Unless required by applicable law or agreed to in writing, software
/// distributed under the License is distributed on an "AS IS" BASIS,
/// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
/// See the License for the specific language governing permissions and
/// limitations under the License.
///
/// Copyright holder is ArangoDB GmbH, Cologne, Germany
///
/// @author Simon Grätzer
////////////////////////////////////////////////////////////////////////////////

#include "IncomingCache.h"
#include "Pregel/CommonFormats.h"
#include "Pregel/Utils.h"
#include "Pregel/WorkerConfig.h"

#include "Basics/MutexLocker.h"
#include "Basics/StaticStrings.h"
#include "Basics/VelocyPackHelper.h"

#include <velocypack/Iterator.h>
#include <velocypack/velocypack-aliases.h>

using namespace arangodb;
using namespace arangodb::pregel;

template <typename M>
InCache<M>::InCache(MessageFormat<M> const* format)
    : _containedMessageCount(0), _format(format) {}

template <typename M>
void InCache<M>::parseMessages(VPackSlice const& incomingData) {
  // every packet contains one shard
  VPackSlice shardSlice = incomingData.get(Utils::shardIdKey);
  VPackSlice messages = incomingData.get(Utils::messagesKey);

  // temporary variables
  VPackValueLength i = 0;
  std::string key;
  prgl_shard_t shard = (prgl_shard_t)shardSlice.getUInt();
  MUTEX_LOCKER(guard, this->_bucketLocker[shard]);

  for (VPackSlice current : VPackArrayIterator(messages)) {
    if (i % 2 == 0) {  // TODO support multiple recipients
      key = current.copyString();
    } else {
      if (current.isArray()) {
        VPackValueLength c = 0;
        for (VPackSlice val : VPackArrayIterator(current)) {
          M newValue;
          _format->unwrapValue(val, newValue);
          _set(shard, key, newValue);
          c++;
        }
        this->_containedMessageCount += c;
      } else {
        M newValue;
        _format->unwrapValue(current, newValue);
        _set(shard, key, newValue);
        this->_containedMessageCount++;
      }
    }
    i++;
  }

  if (i % 2 != 0) {
    THROW_ARANGO_EXCEPTION_MESSAGE(
        TRI_ERROR_BAD_PARAMETER,
        "There must always be a multiple of 2 entries in message array");
  }
}

template <typename M>
void InCache<M>::storeMessageNoLock(prgl_shard_t shard,
                                    std::string const& vertexId,
                                    M const& data) {
  this->_set(shard, vertexId, data);
  this->_containedMessageCount++;
}

template <typename M>
void InCache<M>::storeMessage(prgl_shard_t shard, std::string const& vertexId,
                              M const& data) {
  MUTEX_LOCKER(guard, this->_bucketLocker[shard]);
  this->_set(shard, vertexId, data);
  this->_containedMessageCount++;
}

// ================== ArrayIncomingCache ==================

template <typename M>
ArrayInCache<M>::ArrayInCache(WorkerConfig const* config,
                              MessageFormat<M> const* format)
    : InCache<M>(format) {
  if (config != nullptr) {
    std::set<prgl_shard_t> const& shardIDs = config->localPregelShardIDs();
    // one mutex per shard, we will see how this scales
    for (prgl_shard_t shardID : shardIDs) {
      this->_bucketLocker[shardID];
      _shardMap[shardID];
    }
  }
}

template <typename M>
void ArrayInCache<M>::_set(prgl_shard_t shard, std::string const& key,
                           M const& newValue) {
  HMap& vertexMap(_shardMap[shard]);
  vertexMap[key].push_back(newValue);
}

template <typename M>
void ArrayInCache<M>::mergeCache(WorkerConfig const& config, InCache<M> const* otherCache) {
  ArrayInCache<M>* other = (ArrayInCache<M>*)otherCache;
  this->_containedMessageCount += other->_containedMessageCount;

  // ranomize access to buckets
  std::set<prgl_shard_t> const& shardIDs = config.localPregelShardIDs();
  std::vector<prgl_shard_t> randomized(shardIDs.begin(), shardIDs.end());
  std::random_shuffle(randomized.begin(), randomized.end());
  for (prgl_shard_t shardId : randomized) {
    auto const& it = other->_shardMap.find(shardId);
    if (it != other->_shardMap.end()) {
      MUTEX_LOCKER(guard, this->_bucketLocker[shardId]);
      HMap& myVertexMap = _shardMap[shardId];

      for (auto& vertexMessage : it->second) {
        std::vector<M>& a = myVertexMap[vertexMessage.first];
        std::vector<M> const& b = vertexMessage.second;
        a.insert(a.end(), b.begin(), b.end());
      }
    }
  }
}

template <typename M>
MessageIterator<M> ArrayInCache<M>::getMessages(prgl_shard_t shard,
                                                std::string const& key) {
  HMap const& vertexMap = _shardMap[shard];
  auto vmsg = vertexMap.find(key);
  if (vmsg != vertexMap.end()) {
    // LOG_TOPIC(INFO, Logger::PREGEL) << "Got a message for " << key;
    M const* ptr = vmsg->second.data();
    return MessageIterator<M>(ptr, vmsg->second.size());
  } else {
    // LOG_TOPIC(INFO, Logger::PREGEL) << "No message for " << key;
    return MessageIterator<M>();
  }
}

template <typename M>
void ArrayInCache<M>::clear() {
  for (auto& pair : _shardMap) {
    MUTEX_LOCKER(guard, this->_bucketLocker[pair.first]);
    pair.second.clear();
  }
  this->_containedMessageCount = 0;
}

/// Deletes one entry. DOES NOT LOCK
template <typename M>
void ArrayInCache<M>::erase(prgl_shard_t shard, std::string const& key) {
  HMap& vertexMap = _shardMap[shard];
  auto const& it = vertexMap.find(key);
  if (it != vertexMap.end()) {
    vertexMap.erase(it);
    this->_containedMessageCount--;
  }
}

template <typename M>
void ArrayInCache<M>::forEach(
    std::function<void(prgl_shard_t, std::string const&, M const&)> func) {
  for (auto const& pair : _shardMap) {
    prgl_shard_t shard = pair.first;
    HMap const& vertexMap = pair.second;
    for (auto& vertexMsgs : vertexMap) {
      for (M const& val : vertexMsgs.second) {
        func(shard, vertexMsgs.first, val);
      }
    }
  }
}

// ================== CombiningIncomingCache ==================

template <typename M>
CombiningInCache<M>::CombiningInCache(WorkerConfig const* config,
                                      MessageFormat<M> const* format,
                                      MessageCombiner<M> const* combiner)
    : InCache<M>(format), _combiner(combiner) {
  if (config != nullptr) {
    std::set<prgl_shard_t> const& shardIDs = config->localPregelShardIDs();
    // one mutex per shard, we will see how this scales
    for (prgl_shard_t shardID : shardIDs) {
      this->_bucketLocker[shardID];
      _shardMap[shardID];
    }
  }
}

template <typename M>
void CombiningInCache<M>::_set(prgl_shard_t shard, std::string const& key,
                               M const& newValue) {
  HMap& vertexMap = _shardMap[shard];
  auto vmsg = vertexMap.find(key);
  if (vmsg != vertexMap.end()) {  // got a message for the same vertex
    _combiner->combine(vmsg->second, newValue);
  } else {
    vertexMap.insert(std::make_pair(key, newValue));
  }
}

template <typename M>
void CombiningInCache<M>::mergeCache(WorkerConfig const& config,
                                     InCache<M> const* otherCache) {
  CombiningInCache<M>* other = (CombiningInCache<M>*)otherCache;
  this->_containedMessageCount += other->_containedMessageCount;

  // ranomize access to buckets, don't wait for the lock
  std::set<prgl_shard_t> const& shardIDs = config.localPregelShardIDs();
  std::vector<prgl_shard_t> randomized(shardIDs.begin(), shardIDs.end());
  std::random_shuffle(randomized.begin(), randomized.end());

  size_t i = 0;
  do {
    i = (i + 1) % randomized.size();
    prgl_shard_t shardId = randomized[i];

    auto const& it = other->_shardMap.find(shardId);
    if (it != other->_shardMap.end()) {
      TRY_MUTEX_LOCKER(guard, this->_bucketLocker[shardId]);
      if (guard.isLocked() == false) {
        if (i == 0) {// eventually we hit the last one
          usleep(1000);// don't busy wait
        }
        continue;
      }
      randomized.erase(randomized.begin() + i);

      HMap& myVertexMap = _shardMap[shardId];
      for (auto& vertexMessage : it->second) {
        auto vmsg = myVertexMap.find(vertexMessage.first);
        if (vmsg != myVertexMap.end()) {  // got a message for the same vertex
          _combiner->combine(vmsg->second, vertexMessage.second);
        } else {
          myVertexMap.insert(vertexMessage);
        }
      }
    }
  } while (randomized.size() > 0);
}

template <typename M>
MessageIterator<M> CombiningInCache<M>::getMessages(prgl_shard_t shard,
                                                    std::string const& key) {
  HMap const& vertexMap = _shardMap[shard];
  auto vmsg = vertexMap.find(key);
  if (vmsg != vertexMap.end()) {
    return MessageIterator<M>(&vmsg->second);
  } else {
    return MessageIterator<M>();
  }
}

template <typename M>
void CombiningInCache<M>::clear() {
  for (auto& pair : _shardMap) {
    MUTEX_LOCKER(guard, this->_bucketLocker[pair.first]);
    pair.second.clear();
  }
  this->_containedMessageCount = 0;
}

/// Deletes one entry. DOES NOT LOCK
template <typename M>
void CombiningInCache<M>::erase(prgl_shard_t shard, std::string const& key) {
  HMap& vertexMap = _shardMap[shard];
  auto const& it = vertexMap.find(key);
  if (it != vertexMap.end()) {
    vertexMap.erase(it);
    this->_containedMessageCount--;
  }
}

/// Calls function for each entry. DOES NOT LOCK
template <typename M>
void CombiningInCache<M>::forEach(
    std::function<void(prgl_shard_t shard, std::string const& key, M const&)>
        func) {
  for (auto const& pair : _shardMap) {
    prgl_shard_t shard = pair.first;
    HMap const& vertexMap = pair.second;
    for (auto& vertexMessage : vertexMap) {
      func(shard, vertexMessage.first, vertexMessage.second);
    }
  }
}

// template types to create
template class arangodb::pregel::InCache<int64_t>;
template class arangodb::pregel::InCache<float>;
template class arangodb::pregel::InCache<double>;
template class arangodb::pregel::ArrayInCache<int64_t>;
template class arangodb::pregel::ArrayInCache<float>;
template class arangodb::pregel::ArrayInCache<double>;
template class arangodb::pregel::CombiningInCache<int64_t>;
template class arangodb::pregel::CombiningInCache<float>;
template class arangodb::pregel::CombiningInCache<double>;

// algo specific
template class arangodb::pregel::InCache<SenderMessage<uint64_t>>;
template class arangodb::pregel::ArrayInCache<SenderMessage<uint64_t>>;
template class arangodb::pregel::CombiningInCache<SenderMessage<uint64_t>>;
<commit_msg>Added profiling<commit_after>////////////////////////////////////////////////////////////////////////////////
/// DISCLAIMER
///
/// Copyright 2016 ArangoDB GmbH, Cologne, Germany
///
/// Licensed under the Apache License, Version 2.0 (the "License");
/// you may not use this file except in compliance with the License.
/// You may obtain a copy of the License at
///
///     http://www.apache.org/licenses/LICENSE-2.0
///
/// Unless required by applicable law or agreed to in writing, software
/// distributed under the License is distributed on an "AS IS" BASIS,
/// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
/// See the License for the specific language governing permissions and
/// limitations under the License.
///
/// Copyright holder is ArangoDB GmbH, Cologne, Germany
///
/// @author Simon Grätzer
////////////////////////////////////////////////////////////////////////////////

#include "IncomingCache.h"
#include "Pregel/CommonFormats.h"
#include "Pregel/Utils.h"
#include "Pregel/WorkerConfig.h"

#include "Basics/MutexLocker.h"
#include "Basics/StaticStrings.h"
#include "Basics/VelocyPackHelper.h"

#include <velocypack/Iterator.h>
#include <velocypack/velocypack-aliases.h>

using namespace arangodb;
using namespace arangodb::pregel;

template <typename M>
InCache<M>::InCache(MessageFormat<M> const* format)
    : _containedMessageCount(0), _format(format) {}

template <typename M>
void InCache<M>::parseMessages(VPackSlice const& incomingData) {
  // every packet contains one shard
  VPackSlice shardSlice = incomingData.get(Utils::shardIdKey);
  VPackSlice messages = incomingData.get(Utils::messagesKey);

  // temporary variables
  VPackValueLength i = 0;
  std::string key;
  prgl_shard_t shard = (prgl_shard_t)shardSlice.getUInt();
  MUTEX_LOCKER(guard, this->_bucketLocker[shard]);

  for (VPackSlice current : VPackArrayIterator(messages)) {
    if (i % 2 == 0) {  // TODO support multiple recipients
      key = current.copyString();
    } else {
      if (current.isArray()) {
        VPackValueLength c = 0;
        for (VPackSlice val : VPackArrayIterator(current)) {
          M newValue;
          _format->unwrapValue(val, newValue);
          _set(shard, key, newValue);
          c++;
        }
        this->_containedMessageCount += c;
      } else {
        M newValue;
        _format->unwrapValue(current, newValue);
        _set(shard, key, newValue);
        this->_containedMessageCount++;
      }
    }
    i++;
  }

  if (i % 2 != 0) {
    THROW_ARANGO_EXCEPTION_MESSAGE(
        TRI_ERROR_BAD_PARAMETER,
        "There must always be a multiple of 2 entries in message array");
  }
}

template <typename M>
void InCache<M>::storeMessageNoLock(prgl_shard_t shard,
                                    std::string const& vertexId,
                                    M const& data) {
  this->_set(shard, vertexId, data);
  this->_containedMessageCount++;
}

template <typename M>
void InCache<M>::storeMessage(prgl_shard_t shard, std::string const& vertexId,
                              M const& data) {
  MUTEX_LOCKER(guard, this->_bucketLocker[shard]);
  this->_set(shard, vertexId, data);
  this->_containedMessageCount++;
}

// ================== ArrayIncomingCache ==================

template <typename M>
ArrayInCache<M>::ArrayInCache(WorkerConfig const* config,
                              MessageFormat<M> const* format)
    : InCache<M>(format) {
  if (config != nullptr) {
    std::set<prgl_shard_t> const& shardIDs = config->localPregelShardIDs();
    // one mutex per shard, we will see how this scales
    for (prgl_shard_t shardID : shardIDs) {
      this->_bucketLocker[shardID];
      _shardMap[shardID];
    }
  }
}

template <typename M>
void ArrayInCache<M>::_set(prgl_shard_t shard, std::string const& key,
                           M const& newValue) {
  HMap& vertexMap(_shardMap[shard]);
  vertexMap[key].push_back(newValue);
}

template <typename M>
void ArrayInCache<M>::mergeCache(WorkerConfig const& config, InCache<M> const* otherCache) {
  ArrayInCache<M>* other = (ArrayInCache<M>*)otherCache;
  this->_containedMessageCount += other->_containedMessageCount;

  // ranomize access to buckets
  std::set<prgl_shard_t> const& shardIDs = config.localPregelShardIDs();
  std::vector<prgl_shard_t> randomized(shardIDs.begin(), shardIDs.end());
  std::random_shuffle(randomized.begin(), randomized.end());
  for (prgl_shard_t shardId : randomized) {
    auto const& it = other->_shardMap.find(shardId);
    if (it != other->_shardMap.end()) {
      MUTEX_LOCKER(guard, this->_bucketLocker[shardId]);
      HMap& myVertexMap = _shardMap[shardId];

      for (auto& vertexMessage : it->second) {
        std::vector<M>& a = myVertexMap[vertexMessage.first];
        std::vector<M> const& b = vertexMessage.second;
        a.insert(a.end(), b.begin(), b.end());
      }
    }
  }
}

template <typename M>
MessageIterator<M> ArrayInCache<M>::getMessages(prgl_shard_t shard,
                                                std::string const& key) {
  HMap const& vertexMap = _shardMap[shard];
  auto vmsg = vertexMap.find(key);
  if (vmsg != vertexMap.end()) {
    // LOG_TOPIC(INFO, Logger::PREGEL) << "Got a message for " << key;
    M const* ptr = vmsg->second.data();
    return MessageIterator<M>(ptr, vmsg->second.size());
  } else {
    // LOG_TOPIC(INFO, Logger::PREGEL) << "No message for " << key;
    return MessageIterator<M>();
  }
}

template <typename M>
void ArrayInCache<M>::clear() {
  for (auto& pair : _shardMap) {
    MUTEX_LOCKER(guard, this->_bucketLocker[pair.first]);
    pair.second.clear();
  }
  this->_containedMessageCount = 0;
}

/// Deletes one entry. DOES NOT LOCK
template <typename M>
void ArrayInCache<M>::erase(prgl_shard_t shard, std::string const& key) {
  HMap& vertexMap = _shardMap[shard];
  auto const& it = vertexMap.find(key);
  if (it != vertexMap.end()) {
    vertexMap.erase(it);
    this->_containedMessageCount--;
  }
}

template <typename M>
void ArrayInCache<M>::forEach(
    std::function<void(prgl_shard_t, std::string const&, M const&)> func) {
  for (auto const& pair : _shardMap) {
    prgl_shard_t shard = pair.first;
    HMap const& vertexMap = pair.second;
    for (auto& vertexMsgs : vertexMap) {
      for (M const& val : vertexMsgs.second) {
        func(shard, vertexMsgs.first, val);
      }
    }
  }
}

// ================== CombiningIncomingCache ==================

template <typename M>
CombiningInCache<M>::CombiningInCache(WorkerConfig const* config,
                                      MessageFormat<M> const* format,
                                      MessageCombiner<M> const* combiner)
    : InCache<M>(format), _combiner(combiner) {
  if (config != nullptr) {
    std::set<prgl_shard_t> const& shardIDs = config->localPregelShardIDs();
    // one mutex per shard, we will see how this scales
    for (prgl_shard_t shardID : shardIDs) {
      this->_bucketLocker[shardID];
      _shardMap[shardID];
    }
  }
}

template <typename M>
void CombiningInCache<M>::_set(prgl_shard_t shard, std::string const& key,
                               M const& newValue) {
  HMap& vertexMap = _shardMap[shard];
  auto vmsg = vertexMap.find(key);
  if (vmsg != vertexMap.end()) {  // got a message for the same vertex
    _combiner->combine(vmsg->second, newValue);
  } else {
    vertexMap.insert(std::make_pair(key, newValue));
  }
}

template <typename M>
void CombiningInCache<M>::mergeCache(WorkerConfig const& config,
                                     InCache<M> const* otherCache) {
  CombiningInCache<M>* other = (CombiningInCache<M>*)otherCache;
  this->_containedMessageCount += other->_containedMessageCount;

  // ranomize access to buckets, don't wait for the lock
  std::set<prgl_shard_t> const& shardIDs = config.localPregelShardIDs();
  std::vector<prgl_shard_t> randomized(shardIDs.begin(), shardIDs.end());
  std::random_shuffle(randomized.begin(), randomized.end());

  double t = TRI_microtime();
  size_t i = 0;
  do {
    i = (i + 1) % randomized.size();
    prgl_shard_t shardId = randomized[i];

    auto const& it = other->_shardMap.find(shardId);
    if (it != other->_shardMap.end()) {
      TRY_MUTEX_LOCKER(guard, this->_bucketLocker[shardId]);
      if (guard.isLocked() == false) {
        if (i == 0) {// eventually we hit the last one
          usleep(1000);// don't busy wait
        }
        continue;
      }
      t = TRI_microtime() - t;
      if (t > TRI_SHOW_LOCK_THRESHOLD) {
        LOG_TOPIC(INFO, Logger::PREGEL) << "Getting a bucket took " << t;
      }
      HMap& myVertexMap = _shardMap[shardId];
      for (auto& vertexMessage : it->second) {
        auto vmsg = myVertexMap.find(vertexMessage.first);
        if (vmsg != myVertexMap.end()) {  // got a message for the same vertex
          _combiner->combine(vmsg->second, vertexMessage.second);
        } else {
          myVertexMap.insert(vertexMessage);
        }
      }
      t = TRI_microtime();
    }

    randomized.erase(randomized.begin() + i);
  } while (randomized.size() > 0);
}

template <typename M>
MessageIterator<M> CombiningInCache<M>::getMessages(prgl_shard_t shard,
                                                    std::string const& key) {
  HMap const& vertexMap = _shardMap[shard];
  auto vmsg = vertexMap.find(key);
  if (vmsg != vertexMap.end()) {
    return MessageIterator<M>(&vmsg->second);
  } else {
    return MessageIterator<M>();
  }
}

template <typename M>
void CombiningInCache<M>::clear() {
  for (auto& pair : _shardMap) {
    MUTEX_LOCKER(guard, this->_bucketLocker[pair.first]);
    pair.second.clear();
  }
  this->_containedMessageCount = 0;
}

/// Deletes one entry. DOES NOT LOCK
template <typename M>
void CombiningInCache<M>::erase(prgl_shard_t shard, std::string const& key) {
  HMap& vertexMap = _shardMap[shard];
  auto const& it = vertexMap.find(key);
  if (it != vertexMap.end()) {
    vertexMap.erase(it);
    this->_containedMessageCount--;
  }
}

/// Calls function for each entry. DOES NOT LOCK
template <typename M>
void CombiningInCache<M>::forEach(
    std::function<void(prgl_shard_t shard, std::string const& key, M const&)>
        func) {
  for (auto const& pair : _shardMap) {
    prgl_shard_t shard = pair.first;
    HMap const& vertexMap = pair.second;
    for (auto& vertexMessage : vertexMap) {
      func(shard, vertexMessage.first, vertexMessage.second);
    }
  }
}

// template types to create
template class arangodb::pregel::InCache<int64_t>;
template class arangodb::pregel::InCache<float>;
template class arangodb::pregel::InCache<double>;
template class arangodb::pregel::ArrayInCache<int64_t>;
template class arangodb::pregel::ArrayInCache<float>;
template class arangodb::pregel::ArrayInCache<double>;
template class arangodb::pregel::CombiningInCache<int64_t>;
template class arangodb::pregel::CombiningInCache<float>;
template class arangodb::pregel::CombiningInCache<double>;

// algo specific
template class arangodb::pregel::InCache<SenderMessage<uint64_t>>;
template class arangodb::pregel::ArrayInCache<SenderMessage<uint64_t>>;
template class arangodb::pregel::CombiningInCache<SenderMessage<uint64_t>>;
<|endoftext|>


<commit_before>// Author: Jingyue

#define DEBUG_TYPE "dyn-aa"

#include <string>

#include "llvm/Pass.h"
#include "llvm/Module.h"
#include "llvm/Analysis/AliasAnalysis.h"
#include "llvm/Analysis/Dominators.h"
#include "llvm/Support/raw_ostream.h"
#include "llvm/Transforms/Utils/SSAUpdater.h"

#include "rcs/IntraReach.h"
#include "rcs/IDAssigner.h"

#include "dyn-aa/Utils.h"

using namespace std;
using namespace llvm;
using namespace rcs;
using namespace dyn_aa;

namespace dyn_aa {
struct AliasCheckerInstrumenter: public FunctionPass {
  static const string AssertNoAliasHookName;

  static char ID;

  AliasCheckerInstrumenter();
  virtual void getAnalysisUsage(AnalysisUsage &AU) const;
  virtual bool doInitialization(Module &M);
  virtual bool runOnFunction(Function &F);

 private:
  bool isFree(Function *F) const;
  void addAliasCheckers(Instruction *P, const InstList &Qs);
  void addAliasChecker(Instruction *P, Instruction *Q, SSAUpdater &SU);

  Function *AssertNoAliasHook;
  // Types.
  Type *VoidType;
  IntegerType *CharType, *IntType;
  PointerType *CharStarType;
  // List of freers.
  vector<string> FreeNames;
};
}

static RegisterPass<AliasCheckerInstrumenter> X(
    "instrument-alias-checker",
    "Instrument the alias checker",
    false, // Is CFG Only?
    false); // Is Analysis?

const string AliasCheckerInstrumenter::AssertNoAliasHookName = "AssertNoAlias";

char AliasCheckerInstrumenter::ID = 0;

AliasCheckerInstrumenter::AliasCheckerInstrumenter(): FunctionPass(ID) {
  AssertNoAliasHook = NULL;
  VoidType = NULL;
  CharType = IntType = NULL;
  CharStarType = NULL;
}

void AliasCheckerInstrumenter::getAnalysisUsage(AnalysisUsage &AU) const {
  AU.addRequired<AliasAnalysis>();
  AU.addRequired<IntraReach>();
  AU.addRequired<DominatorTree>();
  AU.addRequired<IDAssigner>();
}

bool AliasCheckerInstrumenter::runOnFunction(Function &F) {
  AliasAnalysis &AA = getAnalysis<AliasAnalysis>();
  IntraReach &IR = getAnalysis<IntraReach>();

  errs() << "Processing function " << F.getName() << "...\n";

  // TODO: consider arguments

  DenseMap<BasicBlock *, InstList> PointerInsts;
  for (Function::iterator BB = F.begin(); BB != F.end(); ++BB) {
    for (BasicBlock::iterator Ins = BB->begin(); Ins != BB->end(); ++Ins) {
      if (!Ins->getType()->isPointerTy())
        continue;
      if (!DynAAUtils::PointerIsAccessed(Ins))
        continue;
      PointerInsts[BB].push_back(Ins);
    }
  }

  // Do not query AA on modified bc. Therefore, we store the checks we are
  // going to add in ToBeChecked, and add them to the program later.
  vector<InstPair> ToBeChecked;
  for (Function::iterator B1 = F.begin(); B1 != F.end(); ++B1) {
    if (!PointerInsts.count(B1))
      continue;
    ConstBBSet ReachableBBs;
    IR.floodfill(B1, ConstBBSet(), ReachableBBs);
    assert(ReachableBBs.count(B1));
    InstList &PointerInstsInB1 = PointerInsts[B1];
    for (size_t i1 = 0, e1 = PointerInstsInB1.size(); i1 < e1; ++i1) {
      Instruction *I1 = PointerInstsInB1[i1];
      for (ConstBBSet::iterator IB2 = ReachableBBs.begin();
           IB2 != ReachableBBs.end(); ++IB2) {
        BasicBlock *B2 = const_cast<BasicBlock *>(*IB2);
        if (!PointerInsts.count(B2))
          continue;
        InstList &PointerInstsInB2 = PointerInsts[B2];
        for (size_t i2 = (B2 == B1 ? i1 + 1 : 0), e2 = PointerInstsInB2.size();
             i2 < e2; ++i2) {
          Instruction *I2 = PointerInstsInB2[i2];
          if (AA.alias(I1, I2) == AliasAnalysis::NoAlias) {
            ToBeChecked.push_back(make_pair(I1, I2));
          }
        }
      }
    }
  }

  errs() << "Adding " << ToBeChecked.size() << " alias checkers...\n";
  for (size_t i = 0; i < ToBeChecked.size(); ) {
    InstList Qs;
    size_t j = i;
    while (j < ToBeChecked.size() &&
           ToBeChecked[j].first == ToBeChecked[i].first) {
      Qs.push_back(ToBeChecked[j].second);
      ++j;
    }
    addAliasCheckers(ToBeChecked[i].first, Qs);
    i = j;
  }

  // Remove deallocators.
  for (Function::iterator BB = F.begin(); BB != F.end(); ++BB) {
    for (BasicBlock::iterator Ins = BB->begin(); Ins != BB->end(); ) {
      BasicBlock::iterator NextIns = Ins; ++NextIns;
      CallSite CS(Ins);
      if (CS && isFree(CS.getCalledFunction()))
        Ins->eraseFromParent();
      Ins = NextIns;
    }
  }

  return true;
}

bool AliasCheckerInstrumenter::doInitialization(Module &M) {
  // Initialize freer list.
  FreeNames.push_back("free");
  FreeNames.push_back("_ZdlPv");
  FreeNames.push_back("_ZdaPv");

  // Initialize basic types.
  VoidType = Type::getVoidTy(M.getContext());
  CharType = Type::getInt8Ty(M.getContext());
  IntType = Type::getInt32Ty(M.getContext());
  CharStarType = PointerType::getUnqual(CharType);

  // Initialize function types.
  vector<Type *> ArgTypes;
  ArgTypes.push_back(CharStarType);
  ArgTypes.push_back(IntType);
  ArgTypes.push_back(CharStarType);
  ArgTypes.push_back(IntType);
  FunctionType *AssertNoAliasHookType = FunctionType::get(VoidType,
                                                          ArgTypes,
                                                          false);

  // Initialize hooks.
  AssertNoAliasHook = Function::Create(AssertNoAliasHookType,
                                       GlobalValue::ExternalLinkage,
                                       AssertNoAliasHookName,
                                       &M);

  return true;
}

void AliasCheckerInstrumenter::addAliasCheckers(Instruction *P,
                                                const InstList &Qs) {
  SSAUpdater SU;
  PointerType *TypeOfP = cast<PointerType>(P->getType());
  SU.Initialize(TypeOfP, P->getName());
  SU.AddAvailableValue(P->getParent(), P);
  if (P->getParent() != P->getParent()->getParent()->begin()) {
    SU.AddAvailableValue(P->getParent()->getParent()->begin(),
                         ConstantPointerNull::get(TypeOfP));
  }

  for (size_t i = 0; i < Qs.size(); ++i) {
    addAliasChecker(P, Qs[i], SU);
  }
}

void AliasCheckerInstrumenter::addAliasChecker(Instruction *P,
                                               Instruction *Q,
                                               SSAUpdater &SU) {
  // It's safe to use DominatorTree here, because SSAUpdater preserves CFG.
  DominatorTree &DT = getAnalysis<DominatorTree>();
  IDAssigner &IDA = getAnalysis<IDAssigner>();

  // Compute the location to add the checker.
  BasicBlock::iterator Loc = Q;
  if (isa<PHINode>(Loc))
    Loc = Loc->getParent()->getFirstNonPHI();
  else if (!Loc->isTerminator())
    ++Loc;
  else
    Loc = cast<InvokeInst>(Loc)->getNormalDest()->getFirstNonPHI();

  // Convert <P> and <Q> to "char *".
  BitCastInst *P2 = new BitCastInst(P, CharStarType, "", Loc);
  BitCastInst *Q2 = new BitCastInst(Q, CharStarType, "", Loc);

  // Compute <P> and <Q>'s value IDs.
  unsigned VIDOfP = IDA.getValueID(P), VIDOfQ = IDA.getValueID(Q);
  assert(VIDOfP != IDAssigner::INVALID_ID && VIDOfQ != IDAssigner::INVALID_ID);

  // Add a function call to AssertNoAlias.
  vector<Value *> Args;
  Args.push_back(P2);
  Args.push_back(ConstantInt::get(IntType, VIDOfP));
  Args.push_back(Q2);
  Args.push_back(ConstantInt::get(IntType, VIDOfQ));
  CallInst::Create(AssertNoAliasHook, Args, "", Loc);

  // The function call just added may be broken, because <P> may not
  // dominate <Q>. Use SSAUpdater to fix it if necessary.
  if (!DT.dominates(P, P2)) {
    assert(P2->getOperand(0) == P);
    SU.RewriteUse(P2->getOperandUse(0));
  }
}

bool AliasCheckerInstrumenter::isFree(Function *F) const {
  if (!F)
    return false;

  vector<string>::const_iterator Pos = find(FreeNames.begin(),
                                            FreeNames.end(),
                                            F->getName());
  return Pos != FreeNames.end();
}
<commit_msg>Enhancement<commit_after>// Author: Jingyue

#define DEBUG_TYPE "dyn-aa"

#include <string>

#include "llvm/Pass.h"
#include "llvm/Module.h"
#include "llvm/Analysis/AliasAnalysis.h"
#include "llvm/Analysis/Dominators.h"
#include "llvm/Support/raw_ostream.h"
#include "llvm/Support/CommandLine.h"
#include "llvm/Transforms/Utils/SSAUpdater.h"

#include "rcs/IntraReach.h"
#include "rcs/IDAssigner.h"

#include "dyn-aa/Utils.h"

using namespace std;
using namespace llvm;
using namespace rcs;
using namespace dyn_aa;

namespace dyn_aa {
struct AliasCheckerInstrumenter: public FunctionPass {
  static const string AssertNoAliasHookName;

  static char ID;

  AliasCheckerInstrumenter();
  virtual void getAnalysisUsage(AnalysisUsage &AU) const;
  virtual bool doInitialization(Module &M);
  virtual bool runOnFunction(Function &F);

 private:
  bool isFree(Function *F) const;
  void addAliasCheckers(Instruction *P, const InstList &Qs);
  void addAliasChecker(Instruction *P, Instruction *Q, SSAUpdater &SU);

  Function *AssertNoAliasHook;
  // Types.
  Type *VoidType;
  IntegerType *CharType, *IntType;
  PointerType *CharStarType;
  // List of freers.
  vector<string> FreeNames;
};
}

static RegisterPass<AliasCheckerInstrumenter> X(
    "instrument-alias-checker",
    "Instrument the alias checker",
    false, // Is CFG Only?
    false); // Is Analysis?

static cl::opt<unsigned> MaxNumAliasChecks(
    "max-alias-checks",
    cl::desc("Add at most this many alias checks. Used for debugging"),
    cl::init((unsigned)-1));

const string AliasCheckerInstrumenter::AssertNoAliasHookName = "AssertNoAlias";

char AliasCheckerInstrumenter::ID = 0;

AliasCheckerInstrumenter::AliasCheckerInstrumenter(): FunctionPass(ID) {
  AssertNoAliasHook = NULL;
  VoidType = NULL;
  CharType = IntType = NULL;
  CharStarType = NULL;
}

void AliasCheckerInstrumenter::getAnalysisUsage(AnalysisUsage &AU) const {
  AU.addRequired<AliasAnalysis>();
  AU.addRequired<IntraReach>();
  AU.addRequired<DominatorTree>();
  AU.addRequired<IDAssigner>();
}

bool AliasCheckerInstrumenter::runOnFunction(Function &F) {
  AliasAnalysis &AA = getAnalysis<AliasAnalysis>();
  IntraReach &IR = getAnalysis<IntraReach>();

  errs() << "Processing function " << F.getName() << "...\n";

  // TODO: consider arguments

  DenseMap<BasicBlock *, InstList> PointerInsts;
  for (Function::iterator BB = F.begin(); BB != F.end(); ++BB) {
    for (BasicBlock::iterator Ins = BB->begin(); Ins != BB->end(); ++Ins) {
      if (!Ins->getType()->isPointerTy())
        continue;
      if (!DynAAUtils::PointerIsAccessed(Ins))
        continue;
      PointerInsts[BB].push_back(Ins);
    }
  }

  // Do not query AA on modified bc. Therefore, we store the checks we are
  // going to add in ToBeChecked, and add them to the program later.
  vector<InstPair> ToBeChecked;
  for (Function::iterator B1 = F.begin(); B1 != F.end(); ++B1) {
    if (!PointerInsts.count(B1))
      continue;
    ConstBBSet ReachableBBs;
    IR.floodfill(B1, ConstBBSet(), ReachableBBs);
    assert(ReachableBBs.count(B1));
    InstList &PointerInstsInB1 = PointerInsts[B1];
    for (size_t i1 = 0, e1 = PointerInstsInB1.size(); i1 < e1; ++i1) {
      Instruction *I1 = PointerInstsInB1[i1];
      for (ConstBBSet::iterator IB2 = ReachableBBs.begin();
           IB2 != ReachableBBs.end(); ++IB2) {
        BasicBlock *B2 = const_cast<BasicBlock *>(*IB2);
        if (!PointerInsts.count(B2))
          continue;
        InstList &PointerInstsInB2 = PointerInsts[B2];
        for (size_t i2 = (B2 == B1 ? i1 + 1 : 0), e2 = PointerInstsInB2.size();
             i2 < e2; ++i2) {
          Instruction *I2 = PointerInstsInB2[i2];
          if (AA.alias(I1, I2) == AliasAnalysis::NoAlias) {
            ToBeChecked.push_back(make_pair(I1, I2));
          }
        }
      }
    }
  }

  unsigned RealSize = ToBeChecked.size();
  if (MaxNumAliasChecks < RealSize)
    RealSize = MaxNumAliasChecks;
  errs() << "Adding " << RealSize << " alias checkers...\n";
  for (size_t i = 0; i < RealSize; ) {
    InstList Qs;
    size_t j = i;
    while (j < RealSize &&
           ToBeChecked[j].first == ToBeChecked[i].first) {
      Qs.push_back(ToBeChecked[j].second);
      ++j;
    }
    addAliasCheckers(ToBeChecked[i].first, Qs);
    i = j;
  }

  // Remove deallocators.
  for (Function::iterator BB = F.begin(); BB != F.end(); ++BB) {
    for (BasicBlock::iterator Ins = BB->begin(); Ins != BB->end(); ) {
      BasicBlock::iterator NextIns = Ins; ++NextIns;
      CallSite CS(Ins);
      if (CS && isFree(CS.getCalledFunction()))
        Ins->eraseFromParent();
      Ins = NextIns;
    }
  }

  return true;
}

bool AliasCheckerInstrumenter::doInitialization(Module &M) {
  // Initialize freer list.
  FreeNames.push_back("free");
  FreeNames.push_back("_ZdlPv");
  FreeNames.push_back("_ZdaPv");

  // Initialize basic types.
  VoidType = Type::getVoidTy(M.getContext());
  CharType = Type::getInt8Ty(M.getContext());
  IntType = Type::getInt32Ty(M.getContext());
  CharStarType = PointerType::getUnqual(CharType);

  // Initialize function types.
  vector<Type *> ArgTypes;
  ArgTypes.push_back(CharStarType);
  ArgTypes.push_back(IntType);
  ArgTypes.push_back(CharStarType);
  ArgTypes.push_back(IntType);
  FunctionType *AssertNoAliasHookType = FunctionType::get(VoidType,
                                                          ArgTypes,
                                                          false);

  // Initialize hooks.
  AssertNoAliasHook = Function::Create(AssertNoAliasHookType,
                                       GlobalValue::ExternalLinkage,
                                       AssertNoAliasHookName,
                                       &M);

  return true;
}

void AliasCheckerInstrumenter::addAliasCheckers(Instruction *P,
                                                const InstList &Qs) {
  SSAUpdater SU;
  PointerType *TypeOfP = cast<PointerType>(P->getType());
  SU.Initialize(TypeOfP, P->getName());
  SU.AddAvailableValue(P->getParent(), P);
  if (P->getParent() != P->getParent()->getParent()->begin()) {
    SU.AddAvailableValue(P->getParent()->getParent()->begin(),
                         ConstantPointerNull::get(TypeOfP));
  }

  for (size_t i = 0; i < Qs.size(); ++i) {
    addAliasChecker(P, Qs[i], SU);
  }
}

void AliasCheckerInstrumenter::addAliasChecker(Instruction *P,
                                               Instruction *Q,
                                               SSAUpdater &SU) {
  // It's safe to use DominatorTree here, because SSAUpdater preserves CFG.
  DominatorTree &DT = getAnalysis<DominatorTree>();
  IDAssigner &IDA = getAnalysis<IDAssigner>();

  // Compute the location to add the checker.
  BasicBlock::iterator Loc = Q;
  if (isa<PHINode>(Loc))
    Loc = Loc->getParent()->getFirstNonPHI();
  else if (!Loc->isTerminator())
    ++Loc;
  else
    Loc = cast<InvokeInst>(Loc)->getNormalDest()->getFirstNonPHI();

  // Convert <P> and <Q> to "char *".
  BitCastInst *P2 = new BitCastInst(P, CharStarType, "", Loc);
  BitCastInst *Q2 = new BitCastInst(Q, CharStarType, "", Loc);

  // Compute <P> and <Q>'s value IDs.
  unsigned VIDOfP = IDA.getValueID(P), VIDOfQ = IDA.getValueID(Q);
  assert(VIDOfP != IDAssigner::INVALID_ID && VIDOfQ != IDAssigner::INVALID_ID);

  // Add a function call to AssertNoAlias.
  vector<Value *> Args;
  Args.push_back(P2);
  Args.push_back(ConstantInt::get(IntType, VIDOfP));
  Args.push_back(Q2);
  Args.push_back(ConstantInt::get(IntType, VIDOfQ));
  CallInst::Create(AssertNoAliasHook, Args, "", Loc);

  // The function call just added may be broken, because <P> may not
  // dominate <Q>. Use SSAUpdater to fix it if necessary.
  if (!DT.dominates(P, P2)) {
    assert(P2->getOperand(0) == P);
    SU.RewriteUse(P2->getOperandUse(0));
  }
}

bool AliasCheckerInstrumenter::isFree(Function *F) const {
  if (!F)
    return false;

  vector<string>::const_iterator Pos = find(FreeNames.begin(),
                                            FreeNames.end(),
                                            F->getName());
  return Pos != FreeNames.end();
}
<|endoftext|>


<commit_before>#include "stdsneezy.h"
#include "database.h"

// create table poll (poll_id int, descr varchar(127), status varchar(8));
// create table poll_option (option_id int, poll_id int, descr varchar(127));
// create table poll_vote (account varchar(80), poll_id int, option_id int);

int ballotBox(TBeing *ch, cmdTypeT cmd, const char *argument, TObj *o, TObj *)
{
  TDatabase db(DB_SNEEZY);
  sstring arg=argument;

  if(cmd==CMD_LIST){
    if(arg.empty()){
      /////////// list all polls
      ch->sendTo(COLOR_BASIC, fmt("You have a look at %s...\n\r") %
		 o->getName());

      db.query("select poll_id, descr from poll where status='open' order by poll_id");

      if(db.fetchRow()){
	ch->sendTo(COLOR_BASIC, "These polls are open:\n\r");
	do {
	  ch->sendTo(COLOR_BASIC, fmt("%-2i| <r>%s<1>\n\r") %
		     convertTo<int>(db["poll_id"]) % db["descr"]);
	} while(db.fetchRow());
	ch->sendTo("\n\r");
      }

      db.query("select poll_id,descr from poll where status='closed' order by poll_id");

      if(db.fetchRow()){
	ch->sendTo(COLOR_BASIC, "These polls are closed:\n\r");
	do {
	  ch->sendTo(COLOR_BASIC, fmt("%-2i| <r>%s<1>\n\r") %
		     convertTo<int>(db["poll_id"]) % db["descr"]);
	} while(db.fetchRow());
	ch->sendTo("\n\r");
      }
    } else {
      //////////// list details of specified poll
      int which=convertTo<int>(arg);
      sstring status="open";

      db.query("select poll_id, descr, status from poll where poll_id=%i",
	       which);

      if(db.fetchRow()){
	ch->sendTo(COLOR_BASIC, fmt("%-2i| <r>%s<1>\n\r") %
		   convertTo<int>(db["poll_id"]) % db["descr"]);
	ch->sendTo(COLOR_BASIC, fmt("Poll is <r>%s<1>.\n\r\n\r") %
		   db["status"]);
	status=db["status"];
      }

      if(status=="open"){
	db.query("select option_id, descr from poll_option where poll_id=%i",
		 which);

	while(db.fetchRow()){
	  ch->sendTo(COLOR_BASIC, fmt("%-2i| <b>%s<1>\n\r") %
		     convertTo<int>(db["option_id"]) % db["descr"]);
	}
      } else if(status=="closed"){
	db.query("select pv.option_id as option_id, po.descr as descr, count(pv.option_id) as count from poll_vote pv, poll_option po where pv.poll_id=%i and pv.poll_id=po.poll_id and pv.option_id=po.option_id group by pv.option_id, po.descr order by count desc", which);

	while(db.fetchRow()){
	  ch->sendTo(COLOR_BASIC, fmt("%-2i| <b>%s<1> (%i votes)\n\r") %
		     convertTo<int>(db["option_id"]) % db["descr"] %
		     convertTo<int>(db["count"]));
	}
      }
    }
    return true;
  } else if(cmd==CMD_VOTE){

    return true;
  }

  return false;
}

int TBeing::doVote(const sstring &)
{
  sendTo(COLOR_BASIC, "There doesn't seem to be a ballot box around.\n\r");
  return true;
}
<commit_msg>fleshed out more - this is the first workable version<commit_after>#include "stdsneezy.h"
#include "database.h"

// create table poll (poll_id int primary key, descr varchar(127), status varchar(8));
// create table poll_option (option_id int, poll_id int, descr varchar(127), primary key (option_id, poll_id));
// create table poll_vote (account varchar(80), poll_id int, option_id int, primary key (account, poll_id, option_id));

int ballotBox(TBeing *ch, cmdTypeT cmd, const char *argument, TObj *o, TObj *)
{
  TDatabase db(DB_SNEEZY);
  sstring arg=argument;

  if(cmd==CMD_LIST){
    if(arg.empty()){
      /////////// list all polls
      ch->sendTo(COLOR_BASIC, fmt("You have a look at %s...\n\r") %
		 o->getName());

      db.query("select poll_id, descr from poll where status='open' order by poll_id");

      if(db.fetchRow()){
	ch->sendTo(COLOR_BASIC, "These polls are open:\n\r");
	do {
	  ch->sendTo(COLOR_BASIC, fmt("%-2i| <r>%s<1>\n\r") %
		     convertTo<int>(db["poll_id"]) % db["descr"]);
	} while(db.fetchRow());
	ch->sendTo("\n\r");
      }

      db.query("select poll_id,descr from poll where status='closed' order by poll_id");

      if(db.fetchRow()){
	ch->sendTo(COLOR_BASIC, "These polls are closed:\n\r");
	do {
	  ch->sendTo(COLOR_BASIC, fmt("%-2i| <r>%s<1>\n\r") %
		     convertTo<int>(db["poll_id"]) % db["descr"]);
	} while(db.fetchRow());
	ch->sendTo("\n\r");
      }
    } else {
      //////////// list details of specified poll
      int which=convertTo<int>(arg);
      sstring status="open";

      db.query("select poll_id, descr, status from poll where poll_id=%i",
	       which);

      if(db.fetchRow()){
	ch->sendTo(COLOR_BASIC, fmt("%-2i| <r>%s<1>\n\r") %
		   convertTo<int>(db["poll_id"]) % db["descr"]);
	ch->sendTo(COLOR_BASIC, fmt("Poll is <r>%s<1>.\n\r\n\r") %
		   db["status"]);
	status=db["status"];
      }

      if(status=="open"){
	db.query("select option_id, descr from poll_option where poll_id=%i",
		 which);

	while(db.fetchRow()){
	  ch->sendTo(COLOR_BASIC, fmt("%-2i| <b>%s<1>\n\r") %
		     convertTo<int>(db["option_id"]) % db["descr"]);
	}
      } else if(status=="closed"){
	db.query("select pv.option_id as option_id, po.descr as descr, count(pv.option_id) as count from poll_vote pv, poll_option po where pv.poll_id=%i and pv.poll_id=po.poll_id and pv.option_id=po.option_id group by pv.option_id, po.descr order by count desc", which);

	while(db.fetchRow()){
	  ch->sendTo(COLOR_BASIC, fmt("%-2i| <b>%s<1> (%i votes)\n\r") %
		     convertTo<int>(db["option_id"]) % db["descr"] %
		     convertTo<int>(db["count"]));
	}
      }
    }
    return true;
  } else if(cmd==CMD_VOTE){
    sstring usage="Usage: vote <poll_id> <option_id>";
    int poll_id, option_id;
    sstring tmp;
    arg=one_argument(arg, tmp);

    if(tmp.empty()){
      if(!ch->isImmortal()){
	ch->sendTo(fmt("%s\n\r") % usage);
      } else {
	ch->sendTo("Usage:\n\r");
	ch->sendTo("vote <poll_id> <option_id>\n\r");
	ch->sendTo("vote <poll_id> add <option_id> <descr>\n\r");
	ch->sendTo("vote <poll_id> remove <option_id>\n\r");
	ch->sendTo("vote add <poll_id> <descr>\n\r");
	ch->sendTo("vote close <poll_id>\n\r");
	ch->sendTo("vote open <poll_id>\n\r");
      }
      return true;
    } if(tmp == "add"){
      // vote add <poll_id> <descr>
      if(!ch->isImmortal()){
	ch->sendTo(fmt("%s\n\r") % usage);
	return true;
      }

      arg=one_argument(arg, tmp);

      db.query("insert into poll values (%i, '%s', 'open')",
	       convertTo<int>(tmp), arg.c_str());
      ch->sendTo("Done.\n\r");
    } else if(tmp == "close"){
      // vote close <poll_id>
      if(!ch->isImmortal()){
	ch->sendTo(fmt("%s\n\r") % usage);
	return true;
      }
      arg=one_argument(arg, tmp);

      db.query("update poll set status='closed' where poll_id=%i",
	       convertTo<int>(tmp));
      ch->sendTo("Done.\n\r");
    } else if(tmp == "open"){
      // vote open <poll_id>
      if(!ch->isImmortal()){
	ch->sendTo(fmt("%s\n\r") % usage);
	return true;
      }
      arg=one_argument(arg, tmp);

      db.query("update poll set status='open' where poll_id=%i",
	       convertTo<int>(tmp));
      ch->sendTo("Done.\n\r");
    } else {
      poll_id=convertTo<int>(tmp);
      arg=one_argument(arg, tmp);

      if(tmp == "add"){
	// vote <poll_id> add <option_id> <descr>
	if(!ch->isImmortal()){
	  ch->sendTo(fmt("%s\n\r") % usage);
	  return true;
	}
	arg=one_argument(arg, tmp);
	option_id=convertTo<int>(tmp);

	db.query("insert into poll_option values (%i, %i, '%s')",
		 option_id, poll_id, arg.c_str());
	ch->sendTo("Done.\n\r");
      } else if(tmp == "remove"){
	// vote <poll_id> remove <option_id>
	if(!ch->isImmortal()){
	  ch->sendTo(fmt("%s\n\r") % usage);
	  return true;
	}
	arg=one_argument(arg, tmp);
	option_id=convertTo<int>(tmp);

	db.query("delete from poll_option where option_id=%i",
		 option_id);
	ch->sendTo("Done.\n\r");
      } else {
	// vote <poll_id> <option_id>
	option_id=convertTo<int>(tmp);

	db.query("select 1 from poll_vote where poll_id=%i and account='%s'",
		 poll_id, ch->desc->account->name);

	if(db.fetchRow()){
	  ch->sendTo("You've already voted in that poll.\n\r");
	  return true;
	}

	db.query("select 1 from poll where poll_id=%i and status='open'",
		 poll_id);

	if(!db.fetchRow()){
	  ch->sendTo("That poll isn't open for voting.\n\r");
	  return true;
	}

	db.query("insert into poll_vote values ('%s', %i, %i)",
		 ch->desc->account->name, poll_id, option_id);

	db.query("select po.descr as descr from poll_option po, poll_vote pv where po.poll_id=%i and po.option_id=%i and pv.poll_id=%i and pv.option_id=%i and pv.account='%s'", poll_id, option_id, poll_id, option_id, ch->desc->account->name);

	if(db.fetchRow()){
	  ch->sendTo(fmt("You cast your vote for %s.\n\r") %
		     db["descr"]);
	} else {
	  ch->sendTo("There was an error - speak to an admin.\n\r");
	}
      }
    }

    return true;
  }

  return false;
}

int TBeing::doVote(const sstring &)
{
  sendTo(COLOR_BASIC, "There doesn't seem to be a ballot box around.\n\r");
  return true;
}
<|endoftext|>