content: /*
   * Licensed to the Apache Software Foundation (ASF) under one
   * or more contributor license agreements.  See the NOTICE file
   * distributed with this work for additional information
   * regarding copyright ownership.  The ASF licenses this file
   * to you 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.
   */
  package org.apache.asterix.lang.sqlpp.visitor.base;
  
  import java.util.Collection;
  import java.util.Collections;
  import java.util.EnumSet;
  import java.util.HashMap;
  import java.util.HashSet;
  import java.util.List;
  import java.util.Map;
  import java.util.Map.Entry;
  import java.util.Set;
  
  import org.apache.asterix.common.exceptions.CompilationException;
  import org.apache.asterix.common.exceptions.ErrorCode;
  import org.apache.asterix.lang.common.base.Expression;
  import org.apache.asterix.lang.common.base.ILangExpression;
  import org.apache.asterix.lang.common.base.IVisitorExtension;
  import org.apache.asterix.lang.common.clause.GroupbyClause;
  import org.apache.asterix.lang.common.clause.LetClause;
  import org.apache.asterix.lang.common.clause.LimitClause;
  import org.apache.asterix.lang.common.context.Scope;
  import org.apache.asterix.lang.common.expression.GbyVariableExpressionPair;
  import org.apache.asterix.lang.common.expression.QuantifiedExpression;
  import org.apache.asterix.lang.common.expression.VariableExpr;
  import org.apache.asterix.lang.common.parser.ScopeChecker;
  import org.apache.asterix.lang.common.rewrites.LangRewritingContext;
  import org.apache.asterix.lang.common.statement.CopyToStatement;
  import org.apache.asterix.lang.common.statement.FunctionDecl;
  import org.apache.asterix.lang.common.statement.InsertStatement;
  import org.apache.asterix.lang.common.statement.Query;
  import org.apache.asterix.lang.common.struct.Identifier;
  import org.apache.asterix.lang.common.struct.QuantifiedPair;
  import org.apache.asterix.lang.common.struct.VarIdentifier;
  import org.apache.asterix.lang.sqlpp.clause.AbstractBinaryCorrelateClause;
  import org.apache.asterix.lang.sqlpp.clause.FromClause;
  import org.apache.asterix.lang.sqlpp.clause.FromTerm;
  import org.apache.asterix.lang.sqlpp.clause.JoinClause;
  import org.apache.asterix.lang.sqlpp.clause.NestClause;
  import org.apache.asterix.lang.sqlpp.clause.SelectSetOperation;
  import org.apache.asterix.lang.sqlpp.clause.UnnestClause;
  import org.apache.asterix.lang.sqlpp.expression.SelectExpression;
  import org.apache.asterix.lang.sqlpp.expression.WindowExpression;
  import org.apache.asterix.lang.sqlpp.struct.SetOperationRight;
  import org.apache.asterix.lang.sqlpp.util.SqlppVariableUtil;
  import org.apache.hyracks.algebricks.common.utils.Pair;
  import org.apache.hyracks.api.exceptions.SourceLocation;
  
  public class AbstractSqlppExpressionScopingVisitor extends AbstractSqlppSimpleExpressionVisitor {
  
      protected final ScopeChecker scopeChecker = new ScopeChecker();
      protected final LangRewritingContext context;
  
      /**
       * @param context, manages ids of variables and guarantees uniqueness of variables.
       */
      public AbstractSqlppExpressionScopingVisitor(LangRewritingContext context) {
          this(context, null);
      }
  
      /**
       * @param context,     manages ids of variables and guarantees uniqueness of variables.
       * @param externalVars pre-defined (external) variables that must be added to the initial scope
       */
      public AbstractSqlppExpressionScopingVisitor(LangRewritingContext context, Collection<VarIdentifier> externalVars) {
          this.context = context;
          this.scopeChecker.setVarCounter(context.getVarCounter());
          if (externalVars != null) {
              for (VarIdentifier paramVar : externalVars) {
                  scopeChecker.getCurrentScope().addSymbolToScope(paramVar);
              }
          }
      }
  
      @Override
      public Expression visit(FromClause fromClause, ILangExpression arg) throws CompilationException {
          Scope scopeForFromClause = scopeChecker.extendCurrentScope();
          for (FromTerm fromTerm : fromClause.getFromTerms()) {
              fromTerm.accept(this, fromClause);
  
              // Merges the variables defined in the current from term into the scope of the current from clause.
              Scope scopeForFromTerm = scopeChecker.removeCurrentScope();
              mergeScopes(scopeForFromClause, scopeForFromTerm, fromTerm.getSourceLocation());
          }
          return null;
      }
  
      @Override
      public Expression visit(FromTerm fromTerm, ILangExpression arg) throws CompilationException {
          // Visit the left expression of a from term.
          fromTerm.setLeftExpression(visit(fromTerm.getLeftExpression(), fromTerm));
  
          scopeChecker.createNewScope();
  
          // Registers the data item variable.
          VariableExpr leftVar = fromTerm.getLeftVariable();
          addNewVarSymbolToScope(scopeChecker.getCurrentScope(), leftVar.getVar(), leftVar.getSourceLocation(),
                  SqlppVariableAnnotation.CONTEXT_VARIABLE);
  
          // Registers the positional variable
          if (fromTerm.hasPositionalVariable()) {
              VariableExpr posVar = fromTerm.getPositionalVariable();
              addNewVarSymbolToScope(scopeChecker.getCurrentScope(), posVar.getVar(), posVar.getSourceLocation());
          }
          // Visits join/unnest/nest clauses.
          for (AbstractBinaryCorrelateClause correlateClause : fromTerm.getCorrelateClauses()) {
              correlateClause.accept(this, fromTerm);
          }
          return null;
      }
  
      @Override
      public Expression visit(JoinClause joinClause, ILangExpression arg) throws CompilationException {
          Scope leftScope = scopeChecker.removeCurrentScope();
          scopeChecker.createNewScope();
          // NOTE: the two join branches cannot be correlated, instead of checking
          // the correlation here,
          // we defer the check to the query optimizer.
          joinClause.setRightExpression(visit(joinClause.getRightExpression(), joinClause));
  
          // Registers the data item variable.
          VariableExpr rightVar = joinClause.getRightVariable();
          addNewVarSymbolToScope(scopeChecker.getCurrentScope(), rightVar.getVar(), rightVar.getSourceLocation(),
                  SqlppVariableAnnotation.CONTEXT_VARIABLE);
  
          if (joinClause.hasPositionalVariable()) {
              // Registers the positional variable.
              VariableExpr posVar = joinClause.getPositionalVariable();
              addNewVarSymbolToScope(scopeChecker.getCurrentScope(), posVar.getVar(), posVar.getSourceLocation());
          }
  
          Scope rightScope = scopeChecker.removeCurrentScope();
          mergeScopes(leftScope, rightScope, joinClause.getRightExpression().getSourceLocation());
          scopeChecker.pushExistingScope(leftScope);
          // The condition expression can refer to the just registered variables
          // for the right branch.
          joinClause.setConditionExpression(visit(joinClause.getConditionExpression(), joinClause));
          return null;
      }
  
      @Override
      public Expression visit(NestClause nestClause, ILangExpression arg) throws CompilationException {
          // NOTE: the two branches of a NEST cannot be correlated, instead of
          // checking the correlation here, we defer the check to the query
          // optimizer.
          nestClause.setRightExpression(visit(nestClause.getRightExpression(), nestClause));
  
          // Registers the data item variable.
          VariableExpr rightVar = nestClause.getRightVariable();
          addNewVarSymbolToScope(scopeChecker.getCurrentScope(), rightVar.getVar(), rightVar.getSourceLocation(),
                  SqlppVariableAnnotation.CONTEXT_VARIABLE);
  
          if (nestClause.hasPositionalVariable()) {
              // Registers the positional variable.
              VariableExpr posVar = nestClause.getPositionalVariable();
              addNewVarSymbolToScope(scopeChecker.getCurrentScope(), posVar.getVar(), posVar.getSourceLocation());
          }
  
          // The condition expression can refer to the just registered variables
          // for the right branch.
          nestClause.setConditionExpression(visit(nestClause.getConditionExpression(), nestClause));
          return null;
      }
  
      @Override
      public Expression visit(UnnestClause unnestClause, ILangExpression arg) throws CompilationException {
          unnestClause.setRightExpression(visit(unnestClause.getRightExpression(), unnestClause));
  
          // register the data item variable
          VariableExpr rightVar = unnestClause.getRightVariable();
          addNewVarSymbolToScope(scopeChecker.getCurrentScope(), rightVar.getVar(), rightVar.getSourceLocation(),
                  SqlppVariableAnnotation.CONTEXT_VARIABLE);
  
          if (unnestClause.hasPositionalVariable()) {
              // register the positional variable
              VariableExpr posVar = unnestClause.getPositionalVariable();
              addNewVarSymbolToScope(scopeChecker.getCurrentScope(), posVar.getVar(), posVar.getSourceLocation());
          }
          return null;
      }
  
      @Override
      public Expression visit(SelectSetOperation selectSetOperation, ILangExpression arg) throws CompilationException {
          Scope scopeBeforeCurrentBranch = scopeChecker.getCurrentScope();
          scopeChecker.createNewScope();
          selectSetOperation.getLeftInput().accept(this, arg);
          if (selectSetOperation.hasRightInputs()) {
              for (SetOperationRight right : selectSetOperation.getRightInputs()) {
                  // Exit scopes that were entered within a previous select expression
                  while (scopeChecker.getCurrentScope() != scopeBeforeCurrentBranch) {
                      scopeChecker.removeCurrentScope();
                  }
                  scopeChecker.createNewScope();
                  right.getSetOperationRightInput().accept(this, arg);
              }
              // Exit scopes that were entered within the last branch of the set operation.
              while (scopeChecker.getCurrentScope() != scopeBeforeCurrentBranch) {
                  scopeChecker.removeCurrentScope();
              }
          }
          return null;
      }
  
      @Override
      public Expression visit(Query q, ILangExpression arg) throws CompilationException {
          q.setBody(visit(q.getBody(), q));
          q.setVarCounter(scopeChecker.getVarCounter());
          return null;
      }
  
      @Override
      public Expression visit(FunctionDecl fd, ILangExpression arg) throws CompilationException {
          scopeChecker.createNewScope();
          fd.setFuncBody(visit(fd.getFuncBody(), fd));
          scopeChecker.removeCurrentScope();
          return null;
      }
  
      @Override
      public Expression visit(GroupbyClause gc, ILangExpression arg) throws CompilationException {
          // After a GROUP BY, variables defined before the current SELECT BLOCK (e.g., in a WITH clause
          // or an outer scope query) should still be visible.
          Scope newScope = new Scope(scopeChecker, scopeChecker.getPrecedingScope());
          // Puts all group-by variables into the symbol set of the new scope.
          Set<VariableExpr> gbyKeyVars = new HashSet<>(); // bindings from prior grouping sets
          for (List<GbyVariableExpressionPair> gbyPairList : gc.getGbyPairList()) {
              for (GbyVariableExpressionPair gbyKeyVarExpr : gbyPairList) {
                  gbyKeyVarExpr.setExpr(visit(gbyKeyVarExpr.getExpr(), gc));
                  VariableExpr gbyKeyVar = gbyKeyVarExpr.getVar();
                  if (gbyKeyVar != null && gbyKeyVars.add(gbyKeyVar)) {
                      addNewVarSymbolToScope(newScope, gbyKeyVar.getVar(), gbyKeyVar.getSourceLocation());
                  }
              }
          }
          if (gc.hasGroupFieldList()) {
              for (Pair<Expression, Identifier> gbyField : gc.getGroupFieldList()) {
                  gbyField.first = visit(gbyField.first, arg);
              }
          }
          if (gc.hasDecorList()) {
              for (GbyVariableExpressionPair decorVarExpr : gc.getDecorPairList()) {
                  decorVarExpr.setExpr(visit(decorVarExpr.getExpr(), gc));
                  VariableExpr decorVar = decorVarExpr.getVar();
                  if (decorVar != null) {
                      addNewVarSymbolToScope(newScope, decorVar.getVar(), decorVar.getSourceLocation());
                  }
              }
          }
          if (gc.hasGroupVar()) {
              VariableExpr groupVar = gc.getGroupVar();
              addNewVarSymbolToScope(newScope, groupVar.getVar(), groupVar.getSourceLocation());
          }
          if (gc.hasWithMap()) {
              Map<Expression, VariableExpr> newWithMap = new HashMap<>();
              for (Entry<Expression, VariableExpr> entry : gc.getWithVarMap().entrySet()) {
                  Expression expr = visit(entry.getKey(), arg);
                  Expression newKey = expr;
                  VariableExpr value = entry.getValue();
                  addNewVarSymbolToScope(newScope, value.getVar(), value.getSourceLocation());
                  newWithMap.put(newKey, value);
              }
              gc.setWithVarMap(newWithMap);
          }
          // Replaces the current scope with the new scope.
          scopeChecker.replaceCurrentScope(newScope);
          return null;
      }
  
      @Override
      public Expression visit(LimitClause limitClause, ILangExpression arg) throws CompilationException {
          scopeChecker.pushForbiddenScope(scopeChecker.getCurrentScope());
          if (limitClause.hasLimitExpr()) {
              limitClause.setLimitExpr(visit(limitClause.getLimitExpr(), limitClause));
          }
          if (limitClause.hasOffset()) {
              limitClause.setOffset(visit(limitClause.getOffset(), limitClause));
          }
          scopeChecker.popForbiddenScope();
          return null;
      }
  
      @Override
      public Expression visit(LetClause letClause, ILangExpression arg) throws CompilationException {
          scopeChecker.extendCurrentScope();
          letClause.setBindingExpr(visit(letClause.getBindingExpr(), letClause));
          VariableExpr varExpr = letClause.getVarExpr();
          addNewVarSymbolToScope(scopeChecker.getCurrentScope(), varExpr.getVar(), varExpr.getSourceLocation());
          return null;
      }
  
      @Override
      public Expression visit(SelectExpression selectExpression, ILangExpression arg) throws CompilationException {
          Scope scopeBeforeSelectExpression = scopeChecker.getCurrentScope();
          scopeChecker.createNewScope();
  
          // visit let list
          if (selectExpression.hasLetClauses()) {
              for (LetClause letClause : selectExpression.getLetList()) {
                  letClause.accept(this, selectExpression);
              }
              scopeChecker.createNewScope();
          }
  
          // visit the main select.
          selectExpression.getSelectSetOperation().accept(this, selectExpression);
  
          // visit order by
          if (selectExpression.hasOrderby()) {
              selectExpression.getOrderbyClause().accept(this, selectExpression);
          }
  
          // visit limit
          if (selectExpression.hasLimit()) {
              selectExpression.getLimitClause().accept(this, selectExpression);
          }
  
          // Exit scopes that were entered within this select expression
          while (scopeChecker.getCurrentScope() != scopeBeforeSelectExpression) {
              scopeChecker.removeCurrentScope();
          }
          return selectExpression;
      }
  
      @Override
      public Expression visit(QuantifiedExpression qe, ILangExpression arg) throws CompilationException {
          scopeChecker.createNewScope();
          for (QuantifiedPair pair : qe.getQuantifiedList()) {
              pair.setExpr(visit(pair.getExpr(), qe));
              VariableExpr varExpr = pair.getVarExpr();
              addNewVarSymbolToScope(scopeChecker.getCurrentScope(), varExpr.getVar(), varExpr.getSourceLocation(),
                      SqlppVariableAnnotation.CONTEXT_VARIABLE);
          }
          qe.setSatisfiesExpr(visit(qe.getSatisfiesExpr(), qe));
          scopeChecker.removeCurrentScope();
          return qe;
      }
  
      @Override
      public Expression visit(VariableExpr varExpr, ILangExpression arg) throws CompilationException {
          String varName = varExpr.getVar().getValue();
          if (scopeChecker.isInForbiddenScopes(varName)) {
              throw new CompilationException(ErrorCode.COMPILATION_ERROR, varExpr.getSourceLocation(),
                      "Inside limit clauses, it is disallowed to reference a variable having the same name as any variable bound in the same scope as the limit clause.");
          }
          Identifier ident = scopeChecker.lookupSymbol(varName);
          if (ident != null) {
              // Exists such an identifier, then this is a variable reference instead of a variable
              // definition.
              varExpr.setIsNewVar(false);
              varExpr.setVar((VarIdentifier) ident);
          }
          return varExpr;
      }
  
      @Override
      public Expression visit(InsertStatement insertStatement, ILangExpression arg) throws CompilationException {
          scopeChecker.createNewScope();
  
          // Visits the body query.
          insertStatement.getQuery().accept(this, insertStatement);
  
          // Registers the (inserted) data item variable.
          VariableExpr bindingVar = insertStatement.getVar();
          if (bindingVar != null) {
              addNewVarSymbolToScope(scopeChecker.getCurrentScope(), bindingVar.getVar(), bindingVar.getSourceLocation(),
                      SqlppVariableAnnotation.CONTEXT_VARIABLE);
          }
  
          // Visits the expression for the returning expression.
          Expression returningExpr = insertStatement.getReturnExpression();
          if (returningExpr != null) {
              insertStatement.setReturnExpression(visit(returningExpr, insertStatement));
          }
          return null;
      }
  
      @Override
      public Expression visit(WindowExpression winExpr, ILangExpression arg) throws CompilationException {
          visitWindowExpressionExcludingExprListAndAggFilter(winExpr, arg);
          if (winExpr.hasWindowVar()) {
              Scope preScope = scopeChecker.getCurrentScope();
              Scope newScope = scopeChecker.extendCurrentScope();
              VariableExpr windowVar = winExpr.getWindowVar();
              addNewVarSymbolToScope(newScope, windowVar.getVar(), windowVar.getSourceLocation());
              winExpr.setExprList(visit(winExpr.getExprList(), arg));
              if (winExpr.hasAggregateFilterExpr()) {
                  winExpr.setAggregateFilterExpr(visit(winExpr.getAggregateFilterExpr(), arg));
              }
              scopeChecker.replaceCurrentScope(preScope);
          } else {
              winExpr.setExprList(visit(winExpr.getExprList(), arg));
              if (winExpr.hasAggregateFilterExpr()) {
                  winExpr.setAggregateFilterExpr(visit(winExpr.getAggregateFilterExpr(), arg));
              }
          }
          return winExpr;
      }
  
      @Override
      public Expression visit(CopyToStatement stmtCopy, ILangExpression arg) throws CompilationException {
          stmtCopy.setBody(stmtCopy.getBody().accept(this, stmtCopy));
  
          // Scope that only contains the source variable
          Scope currentNewScope = scopeChecker.createNewScope();
          VariableExpr sourceVarExpr = stmtCopy.getSourceVariable();
          addNewVarSymbolToScope(currentNewScope, sourceVarExpr.getVar(), stmtCopy.getSourceLocation());
  
          // Visit partition exprs
          stmtCopy.setPartitionExpressions(visit(stmtCopy.getPartitionExpressions(), stmtCopy));
  
          // Add partition variables to the current scope
          Collection<VariableExpr> partitionVars = stmtCopy.getPartitionsVariables().values();
          for (VariableExpr partVarExpr : partitionVars) {
              addNewVarSymbolToScope(currentNewScope, partVarExpr.getVar(), partVarExpr.getSourceLocation());
          }
  
          // Visit order by
          stmtCopy.setOrderByList(visit(stmtCopy.getOrderByList(), stmtCopy));
  
          // Visit path exprs
          stmtCopy.setPathExpressions(visit(stmtCopy.getPathExpressions(), stmtCopy));
  
          // Visit key exprs
          stmtCopy.setKeyExpressions(visit(stmtCopy.getKeyExpressions(), stmtCopy));
  
          return null;
      }
  
      @Override
      public Expression visit(IVisitorExtension ve, ILangExpression arg) throws CompilationException {
          return ve.variableScopeDispatch(this, arg, scopeChecker);
      }
  
      // Adds a new encountered alias identifier into a scope
      private void addNewVarSymbolToScope(Scope scope, VarIdentifier var, SourceLocation sourceLoc,
              SqlppVariableAnnotation... varAnnotations) throws CompilationException {
          if (scope.findLocalSymbol(var.getValue()) != null) {
              throw new CompilationException(ErrorCode.COMPILATION_ERROR, sourceLoc,
                      "Duplicate alias definitions: " + SqlppVariableUtil.toUserDefinedName(var.getValue()));
          }
          Set<SqlppVariableAnnotation> annotations;
          if (varAnnotations == null || varAnnotations.length == 0) {
              annotations = Collections.emptySet();
          } else {
              annotations = EnumSet.noneOf(SqlppVariableAnnotation.class);
              Collections.addAll(annotations, varAnnotations);
          }
          scope.addNewVarSymbolToScope(var, annotations);
      }
  
      // Merges <code>scopeToBeMerged</code> into <code>hostScope</code>.
      protected void mergeScopes(Scope hostScope, Scope scopeToBeMerged, SourceLocation sourceLoc)
              throws CompilationException {
          Set<String> symbolsToBeMerged = scopeToBeMerged.getLocalSymbols();
          for (String symbolToBeMerged : symbolsToBeMerged) {
              if (hostScope.findLocalSymbol(symbolToBeMerged) != null) {
                  throw new CompilationException(ErrorCode.COMPILATION_ERROR, sourceLoc,
                          "Duplicate alias definitions: " + SqlppVariableUtil.toUserDefinedName(symbolToBeMerged));
              }
          }
          hostScope.merge(scopeToBeMerged);
      }
  
      public enum SqlppVariableAnnotation implements Scope.SymbolAnnotation {
          /**
           * Context variables are those that participate in the second stage of the name resolution process.
           * A single name identifier is first attempted to be resolved as a variable reference. If that fails
           * because there's no variable with such name then (second stage) it's resolved as a field access on a context
           * variable (if there's only one context variable defined in the local scope).
           * <p>
           * See {@link org.apache.asterix.lang.sqlpp.rewrites.visitor.VariableCheckAndRewriteVisitor}
           */
          CONTEXT_VARIABLE
      }
  }
  