text: #include "lower_index_expressions.h"
  
  #include "ir.h"
  #include "ir_rewriter.h"
  #include "ir_transforms.h"
  
  #include "lower_indexexprs.h"
  #include "lower_scatter_workspace.h"
  #include "lower_transpose.h"
  #include "lower_matrix_multiply.h"
  
  #include "path_expressions.h"
  
  namespace simit {
  namespace ir {
  
  inline unsigned getExperssionArity(const IndexExpr* iexpr) {
    // The expression arity is the number of binary expressions plus one
    unsigned expressionArity = 1;
    match(iexpr->value,
      std::function<void(const BinaryExpr*)>([&](const BinaryExpr* op) {
        ++expressionArity;
      })
    );
    return expressionArity;
  }
  
  inline bool isBinaryExpression(const IndexExpr* iexpr) {
    return getExperssionArity(iexpr) == 2;
  }
  
  inline bool isScale(const IndexExpr* iexpr) {
  
    bool result = false;
    // Expression with at least one scalar variable
    match(iexpr->value,
      std::function<void(const VarExpr*)>([&](const VarExpr* op) {
        if (isScalar(op->type)) {
          result = true;
        }
      }),
      std::function<void(const Literal*)>([&](const Literal* op) {
        if (isScalar(op->type)) {
          result = true;
        }
      }),
      std::function<void(const Load*)>([&](const Load* op) {
        if (isScalar(op->type)) {
          result = true;
        }
      }),
      std::function<void(const FieldRead*)>([&](const FieldRead* op) {
        if (isScalar(op->type)) {
          result = true;
        }
      })
    );
  
    return result;
  }
  
  inline bool doesOperandsHaveSameStructureOrIsDiagonal(const IndexExpr* iexpr,
                                                        const Storage& storage) {
    bool result = true;
    pe::PathExpression pexpr;
    match(iexpr->value,
      std::function<void(const VarExpr*)>([&](const VarExpr* op) {
        simit_iassert(storage.hasStorage(op->var));
        auto tensorStorage = storage.getStorage(op->var);
      
        if (tensorStorage.getKind() != TensorStorage::Diagonal) {
          auto p = storage.getStorage(op->var).getTensorIndex().getPathExpression();
          if (!pexpr.defined()) {
            pexpr = p;
          }
          else if (pexpr != p) {
            result = false;
          }
        }
      })
    );
  
    return result;
  }
  
  inline bool isElwise(const IndexExpr* iexpr) {
    bool result = true;
    match(iexpr->value,
      std::function<void(const IndexedTensor*)>([&](const IndexedTensor* op) {
        if (result == false) {
          return;
        }
        if (op->indexVars.size() != iexpr->resultVars.size() ||
            !std::equal(op->indexVars.begin(), op->indexVars.end(),
                        iexpr->resultVars.begin())) {
          result = false;
          return;
        }
      })
    );
  
    return result;
  }
  
  inline bool isTranspose(const IndexExpr* iexpr) {
    auto outvars  = iexpr->resultVars;
    if (outvars.size() != 2) {
      return false;
    }
    if (!isa<IndexedTensor>(iexpr->value)) {
      return false;
    }
    auto invars = to<IndexedTensor>(iexpr->value)->indexVars;
    if (invars.size() !=2) {
      return false;
    }
    if (outvars[0] != invars[1] || outvars[1] != invars[0]) {
      return false;
    }
    return true;
  }
  
  inline bool isGemm(const IndexExpr* iexpr) {
    // Very specific index expression form: (i,j B(i,+k)*C(+k,j))
    // "First" matrix is defined as the one with its first index var
    // as a free variable. The "second" matrix is defined as the reverse.
    // The remaining variable should be summed.
    if (iexpr->resultVars.size() != 2) {
      return false;
    }
    bool result = true;
    bool foundFirst = false;
    bool foundSecond = false;
    match(iexpr->value,
      std::function<void(const IndexedTensor*)>([&](const IndexedTensor* op) {
        if (result == false) {
          return;
        }
        if (op->indexVars.size() != 2) {
          result = false;
          return;
        }
        // Check for first matrix
        // TODO: We only allow non-transposed multiplication
        if (op->indexVars[0].isFreeVar() &&
            op->indexVars[0] == iexpr->resultVars[0] &&
            op->indexVars[1].isReductionVar() &&
            op->indexVars[1].getOperator() == ReductionOperator::Sum) {
          // We're doing a transpose multiply, unhandled right now
          if (foundFirst) {
            result = false;
            return;
          }
          foundFirst = true;
        }
        // Check for second matrix
        else if (op->indexVars[1].isFreeVar() &&
                 op->indexVars[1] == iexpr->resultVars[1] &&
                 op->indexVars[0].isReductionVar() &&
                 op->indexVars[0].getOperator() == ReductionOperator::Sum) {
          // Transpose multiply is unhandled right now
          if (foundSecond) {
            result = false;
            return;
          }
          foundSecond = true;
        }
        // Tensor term is not of matrix multiply form
        else {
          result = false;
          return;
        }
      })
    );
    result = result && foundFirst && foundSecond;
  
    return result;
  }
  
  Func lowerIndexExpressions(Func func) {
    class LowerIndexExpressionsRewriter : private IRRewriter {
    public:
      Func lower(Func func) {
        storage = &func.getStorage();
        environment = func.getEnvironment();
        return this->rewrite(func);
      }
  
    private:
      Storage *storage;
      Environment environment;
      
      using IRRewriter::visit;
  
      void visit(const Func* f) {
        Stmt body = rewrite(f->getBody());
        if (body != f->getBody()) {
          func = Func(f->getName(), f->getArguments(), f->getResults(),
                      body, environment, f->getKind());
          func.setStorage(*storage);
        }
        else {
          func = *f;
        }
      }
  
      void visit(const AssignStmt *op) {
        if (!isa<IndexExpr>(op->value) && op->cop == CompoundOperator::None) {
          IRRewriter::visit(op);
          return;
        }
        if (!isa<IndexExpr>(op->value)) {
          stmt = lowerIndexStatement(op, &environment, *storage);
          return;
        }
  
        const IndexExpr* iexpr = to<IndexExpr>(op->value);
  
        // Dispatch the index expression lowering to the correct lowering pass.
        enum Kind {Unknown, DenseResult, MatrixScale,
                   MatrixElwiseWithSameStructureOrDiagonal, MatrixElwise,
                   MatrixTranspose, MatrixMultiply};
        Kind kind = Unknown;
  
        simit_iassert(iexpr->type.isTensor());
        const Var& var = op->var;
        const TensorType* type = iexpr->type.toTensor();
  
        if (type->order()==0 || type->order()==1 ||
            storage->getStorage(var).getKind() == TensorStorage::Dense) {
          kind = DenseResult;
        }
        else if (isScale(iexpr)) {
          kind = MatrixScale;
        }
        else if (isElwise(iexpr)) {
          kind = doesOperandsHaveSameStructureOrIsDiagonal(iexpr, *storage)
                 ? MatrixElwiseWithSameStructureOrDiagonal
                 : MatrixElwise;
        }
        else if (isTranspose(iexpr)) {
          kind = MatrixTranspose;
        }
        else if (isGemm(iexpr)) {
          kind = MatrixMultiply;
        }
  
        simit_iassert(kind != Unknown)
            << "Index expression lowering does not know how to lower: "
            << Stmt(op);
  
        switch (kind) {
          case DenseResult:
          case MatrixScale:
          case MatrixElwiseWithSameStructureOrDiagonal:
            stmt = lowerIndexStatement(op, &environment, *storage);
            break;
          case MatrixElwise:
            stmt = lowerScatterWorkspace(op->var, iexpr, &environment, storage);
            break;
          case MatrixTranspose:
            stmt = lowerTranspose(op->var, iexpr, &environment, storage);
            break;
          case MatrixMultiply:
            stmt = lowerMatrixMultiply(op->var, iexpr, &environment, storage);
            break;
          case Unknown:
            simit_unreachable << "unknown matrix expression";
            break;
        }
        simit_iassert(stmt.defined());
        stmt = Comment::make(util::toString(*op), stmt, false, true);
      }
  
      void visit(const FieldWrite *op) {
        if (!isa<IndexExpr>(op->value) && op->cop == CompoundOperator::None) {
          IRRewriter::visit(op);
          return;
        }
        stmt = lowerIndexStatement(op, &environment, *storage);
  
        if (isa<IndexExpr>(op->value)) {
          stmt = Comment::make(util::toString(*op), stmt, false, true);
        }
      }
  
      void visit(const TensorWrite *op) {
        if (!isa<IndexExpr>(op->value) && op->cop == CompoundOperator::None) {
          IRRewriter::visit(op);
          return;
        }
        stmt = lowerIndexStatement(op, &environment, *storage);
        if (isa<IndexExpr>(op->value)) {
          stmt = Comment::make(util::toString(*op), stmt, false, true);
        }
      }
  
      void visit(const IndexExpr *op) {
        simit_iassert_scalar(Expr(op));
        expr = rewrite(op->value);
      }
    };
    func = LowerIndexExpressionsRewriter().lower(func);
    func = insertVarDecls(func);
  
    return func;
  }
  
  }}
  