njs_module_hash_test(njs_lvlhsh_query_t *lhq, void *data)
{
    njs_mod_t  *module;

    module = data;

    if (njs_strstr_eq(&lhq->key, &module->name)) {
        return NJS_OK;
    }

    return NJS_DECLINED;
}


njs_module_lookup(njs_vm_t *vm, const njs_str_t *cwd, njs_module_info_t *info)
{
    njs_int_t   ret;
    njs_str_t   *path;
    njs_uint_t  i;

    if (info->name.start[0] == '/') {
        return njs_module_path(vm, NULL, info);
    }

    ret = njs_module_path(vm, cwd, info);

    if (ret != NJS_DECLINED) {
        return ret;
    }

    if (vm->paths == NULL) {
        return NJS_DECLINED;
    }

    path = vm->paths->start;

    for (i = 0; i < vm->paths->items; i++) {
        ret = njs_module_path(vm, path, info);

        if (ret != NJS_DECLINED) {
            return ret;
        }

        path++;
    }

    return NJS_DECLINED;
}

Fixed typo while calculating module path length.

The issue was introduced in 77c398f26d7e (not released yet).


njs_default_module_loader(njs_vm_t *vm, njs_external_ptr_t external,
    njs_str_t *name)
{
    njs_int_t          ret;
    njs_str_t          cwd, text;
    njs_parser_t       *prev;
    njs_mod_t          *module;
    njs_module_info_t  info;

    prev = external;

    njs_memzero(&info, sizeof(njs_module_info_t));

    info.name = *name;
    njs_file_dirname(&prev->lexer->file, &cwd);

    ret = njs_module_lookup(vm, &cwd, &info);
    if (njs_slow_path(ret != NJS_OK)) {
        return NULL;
    }

    ret = njs_module_read(vm, info.fd, &text);

    (void) close(info.fd);

    if (njs_slow_path(ret != NJS_OK)) {
        njs_internal_error(vm, "while reading \"%V\" module", &info.file);
        return NULL;
    }

    module = njs_vm_compile_module(vm, &info.file, &text.start,
                                   &text.start[text.length]);

    njs_mp_free(vm->mem_pool, text.start);

    return module;
}


get_user_rules(svn_authz_t *authz,
               const char *repos_name,
               const char *user)
{
  apr_pool_t *pool;

  /* Search our cache for a suitable previously filtered tree. */
  if (authz->filtered)
    {
      /* Is this a suitable filtered tree? */
      if (matches_filtered_tree(authz->filtered, repos_name, user))
        return authz->filtered;

      /* Drop the old filtered tree before creating a new one. */
      svn_pool_destroy(authz->filtered->pool);
      authz->filtered = NULL;
    }

  /* Global cache lookup.  Filter the full model only if necessary. */
  pool = svn_pool_create(authz->pool);

  /* Write a new entry. */
  authz->filtered = apr_palloc(pool, sizeof(*authz->filtered));
  authz->filtered->pool = pool;
  authz->filtered->repository = apr_pstrdup(pool, repos_name);
  authz->filtered->user = user ? apr_pstrdup(pool, user) : NULL;
  authz->filtered->lookup_state = create_lookup_state(pool);
  authz->filtered->root = NULL;

  svn_authz__get_global_rights(&authz->filtered->global_rights,
                               authz->full, user, repos_name);

  return authz->filtered;
}

Fix issue #4880 "Use-after-free of object-pools when used as httpd module"

Ensure that we initialize authz again if the pool which our authz
caches depend on is cleared. Apache HTTPD may run pre/post config
hooks multiple times and clear its global configuration pool which
our authz caching pools depend on.

Reported-by: Thomas Weißschuh (thomas {at} t-8ch dot de)

Thomas has also confirmed that this patch fixes the problem.

* subversion/libsvn_repos/authz.c
  (deinit_authz): New pool cleanup handler which resets authz initialization
   in case the parent pool of our authz caches is cleared.
  (synchronized_authz_initialize): Register new pool cleanup handler.

git-svn-id: https://svn.apache.org/repos/asf/subversion/trunk@1894734 13f79535-47bb-0310-9956-ffa450edef68


add_complex_matches(lookup_state_t *state,
                    const svn_stringbuf_t *segment,
                    apr_array_header_t *patterns)
{
  int i;
  for (i = 0; i < patterns->nelts; ++i)
    {
      node_t *node = APR_ARRAY_IDX(patterns, i, sorted_pattern_t).node;
      if (0 == apr_fnmatch(node->segment.data, segment->data, 0))
        add_next_node(state, node);
    }
}


finalize_subnode_hash(apr_hash_t *hash,
                      limited_rights_t *sum,
                      apr_pool_t *scratch_pool)
{
  if (hash)
    {
      apr_hash_index_t *hi;
      for (hi = apr_hash_first(scratch_pool, hash);
           hi;
           hi = apr_hash_next(hi))
        finalize_tree(apr_hash_this_val(hi), sum, scratch_pool);
    }
}


ensure_pattern_sub_nodes(node_t *node,
                         apr_pool_t *result_pool)
{
  if (node->pattern_sub_nodes == NULL)
    node->pattern_sub_nodes = apr_pcalloc(result_pool,
                                          sizeof(*node->pattern_sub_nodes));

  return node->pattern_sub_nodes;
}


lookup(lookup_state_t *state,
       const char *path,
       authz_access_t required,
       svn_boolean_t recursive,
       apr_pool_t *scratch_pool)
{
  /* Create a scratch pad large enough to hold any of PATH's segments. */
  apr_size_t path_len = strlen(path);
  svn_stringbuf_ensure(state->scratch_pad, path_len);

  /* Normalize start and end of PATH.  Most paths will be fully normalized,
   * so keep the overhead as low as possible. */
  if (path_len && path[path_len-1] == '/')
    {
      do
      {
        --path_len;
      }
      while (path_len && path[path_len-1] == '/');
      path = apr_pstrmemdup(scratch_pool, path, path_len);
    }

  while (path[0] == '/')
    ++path;     /* Don't update PATH_LEN as we won't need it anymore. */

  /* Actually walk the path rule tree following PATH until we run out of
   * either tree or PATH. */
  while (state->current->nelts && path)
    {
      apr_array_header_t *temp;
      int i;
      svn_stringbuf_t *segment = state->scratch_pad;

      /* Shortcut 1: We could nowhere find enough rights in this sub-tree. */
      if ((state->rights.max_rights & required) != required)
        return FALSE;

      /* Shortcut 2: We will find enough rights everywhere in this sub-tree. */
      if ((state->rights.min_rights & required) == required)
        return TRUE;

      /* Extract the next segment. */
      path = next_segment(segment, path);

      /* Initial state for this segment. */
      apr_array_clear(state->next);
      state->rights.access.sequence_number = NO_SEQUENCE_NUMBER;
      state->rights.access.rights = authz_access_none;

      /* These init values ensure that the first node's value will be used
       * when combined with them.  If there is no first node,
       * state->access.sequence_number remains unchanged and we will use
       * the parent's (i.e. inherited) access rights. */
      state->rights.min_rights = authz_access_write;
      state->rights.max_rights = authz_access_none;

      /* Update the PARENT_PATH member in STATE to match the nodes in
       * CURRENT at the end of this iteration, i.e. if and when NEXT
       * has become CURRENT. */
      if (path)
        {
          svn_stringbuf_appendbyte(state->parent_path, '/');
          svn_stringbuf_appendbytes(state->parent_path, segment->data,
                                    segment->len);
        }

      /* Scan follow all alternative routes to the next level. */
      for (i = 0; i < state->current->nelts; ++i)
        {
          node_t *node = APR_ARRAY_IDX(state->current, i, node_t *);
          if (node->sub_nodes)
            add_next_node(state, apr_hash_get(node->sub_nodes, segment->data,
                                              segment->len));

          /* Process alternative, wildcard-based sub-nodes. */
          if (node->pattern_sub_nodes)
            {
              add_next_node(state, node->pattern_sub_nodes->any);

              /* If the current node represents a "**" pattern, it matches
               * to all levels. So, add it to the list for the NEXT level. */
              if (node->pattern_sub_nodes->repeat)
                add_next_node(state, node);

              /* Find all prefix pattern matches. */
              if (node->pattern_sub_nodes->prefixes)
                add_prefix_matches(state, segment,
                                   node->pattern_sub_nodes->prefixes);

              if (node->pattern_sub_nodes->complex)
                add_complex_matches(state, segment,
                                    node->pattern_sub_nodes->complex);

              /* Find all suffux pattern matches.
               * This must be the last check as it destroys SEGMENT. */
              if (node->pattern_sub_nodes->suffixes)
                {
                  /* Suffixes behave like reversed prefixes. */
                  svn_authz__reverse_string(segment->data, segment->len);
                  add_prefix_matches(state, segment,
                                     node->pattern_sub_nodes->suffixes);
                }
            }
        }

      /* If no rule applied to this SEGMENT directly, the parent rights
       * will apply to at least the SEGMENT node itself and possibly
       * other parts deeper in it's subtree. */
      if (!has_local_rule(&state->rights))
        {
          state->rights.access = state->parent_rights.access;
          state->rights.min_rights &= state->parent_rights.access.rights;
          state->rights.max_rights |= state->parent_rights.access.rights;
        }

      /* The list of nodes for SEGMENT is now complete.  If we need to
       * continue, make it the current and put the old one into the recycler.
       *
       * If this is the end of the path, keep the parent path and rights in
       * STATE as are such that sibling lookups will benefit from it.
       */
      if (path)
        {
          temp = state->current;
          state->current = state->next;
          state->next = temp;

          /* In STATE, PARENT_PATH, PARENT_RIGHTS and CURRENT are now in sync. */
          state->parent_rights = state->rights;
        }
    }

  /* If we check recursively, none of the (potential) sub-paths must have
   * less than the REQUIRED access rights.  "Potential" because we don't
   * verify that the respective paths actually exist in the repository.
   */
  if (recursive)
    return (state->rights.min_rights & required) == required;

  /* Return whether the access rights on PATH fully include REQUIRED. */
  return (state->rights.access.rights & required) == required;
}


add_next_node(lookup_state_t *state,
              node_t *node)
{
  /* Allowing NULL nodes simplifies the caller. */
  if (node)
    {
      /* The rule with the highest sequence number is the one that applies.
       * Not all nodes that we are following have rules that apply directly
       * to this path but are mere intermediates that may only have some
       * matching deep sub-node. */
      combine_access(&state->rights, &node->rights);

      /* The rule tree node can be seen as an overlay of all the nodes that
       * we are following.  Any of them _may_ match eventually, so the min/
       * max possible access rights are a combination of all these sub-trees.
       */
      combine_right_limits(&state->rights, &node->rights);

      /* NODE is now enlisted as a (potential) match for the next segment. */
      APR_ARRAY_PUSH(state->next, node_t *) = node;

      /* Variable length sub-segment sequences apply to the same node as
       * they match empty sequences as well. */
      if (node->pattern_sub_nodes && node->pattern_sub_nodes->any_var)
        {
          node = node->pattern_sub_nodes->any_var;

          /* This is non-recursive due to ACL normalization. */
          combine_access(&state->rights, &node->rights);
          combine_right_limits(&state->rights, &node->rights);
          APR_ARRAY_PUSH(state->next, node_t *) = node;
        }
    }
}


has_local_rule(const limited_rights_t *rights)
{
  return rights->access.sequence_number != NO_SEQUENCE_NUMBER;
}


link_prefix_patterns(apr_array_header_t *array)
{
  int i;
  if (!array)
    return;

  for (i = 1; i < array->nelts; ++i)
    {
      sorted_pattern_t *prev
        = &APR_ARRAY_IDX(array, i - 1, sorted_pattern_t);
      sorted_pattern_t *pattern
        = &APR_ARRAY_IDX(array, i, sorted_pattern_t);

      /* Does PATTERN potentially have a prefix in ARRAY?
       * If so, at least the first char must match with the predecessor's
       * because the array is sorted by that string. */
      if (prev->node->segment.data[0] != pattern->node->segment.data[0])
        continue;

      /* Only the predecessor or any of its prefixes can be the closest
       * prefix to PATTERN. */
      for ( ; prev; prev = prev->next)
        if (   prev->node->segment.len < pattern->node->segment.len
            && !memcmp(prev->node->segment.data,
                       pattern->node->segment.data,
                       prev->node->segment.len))
          {
            pattern->next = prev;
            break;
          }
    }
}


next_segment(svn_stringbuf_t *segment,
             const char *path)
{
  apr_size_t len;
  char c;

  /* Read and scan PATH for NUL and '/' -- whichever comes first. */
  for (len = 0, c = *path; c; c = path[++len])
    if (c == '/')
      {
        /* End of segment. */
        segment->data[len] = 0;
        segment->len = len;

        /* If PATH is not normalized, this is where we skip whole sequences
         * of separators. */
        while (path[++len] == '/')
          ;

        /* Continue behind the last separator in the sequence.  We will
         * treat trailing '/' as indicating an empty trailing segment.
         * Therefore, we never have to return NULL here. */
        return path + len;
      }
    else
      {
        /* Copy segment contents directly into the result buffer.
         * On many architectures, this is almost or entirely for free. */
        segment->data[len] = c;
      }

  /* No separator found, so all of PATH has been the last segment. */
  segment->data[len] = 0;
  segment->len = len;

  /* Tell the caller that this has been the last segment. */
  return NULL;
}


compare_node_path_segment(const void *void_lhs,
                          const void *void_rhs)
{
  const sorted_pattern_t *element = void_lhs;
  const char *segment = void_rhs;

  return strcmp(element->node->segment.data, segment);
}


create_node(authz_rule_segment_t *segment,
            apr_pool_t *result_pool)
{
  node_t *result = apr_pcalloc(result_pool, sizeof(*result));
  if (segment)
    result->segment = segment->pattern;
  else
    {
      result->segment.data = "";
      result->segment.len = 0;
    }
  result->rights.access.sequence_number = NO_SEQUENCE_NUMBER;
  return result;
}


ensure_node(node_t **node,
            authz_rule_segment_t *segment,
            apr_pool_t *result_pool)
{
  if (!*node)
    *node = create_node(segment, result_pool);

  return *node;
}


matches_filtered_tree(const authz_user_rules_t *authz,
                      const char *repos_name,
                      const char *user)
{
  /* Does the user match? */
  if (user)
    {
      if (authz->user == NULL || strcmp(user, authz->user))
        return FALSE;
    }
  else if (authz->user != NULL)
    return FALSE;

  /* Does the repository match as well? */
  return strcmp(repos_name, authz->repository) == 0;
}


construct_filtered_key(const char *repos_name,
                       const char *user,
                       const svn_membuf_t *authz_id,
                       apr_pool_t *result_pool)
{
  svn_membuf_t *result = apr_pcalloc(result_pool, sizeof(*result));
  size_t repos_len = strlen(repos_name);
  size_t user_len = user ? strlen(user) : 1;
  const char *nullable_user = user ? user : "\0";
  size_t size = authz_id->size + repos_len + 1 + user_len + 1;

  svn_membuf__create(result, size, result_pool);
  result->size = size;

  memcpy(result->data, repos_name, repos_len + 1);
  size = repos_len + 1;
  memcpy((char *)result->data + size, nullable_user, user_len + 1);
  size += user_len + 1;
  memcpy((char *)result->data + size, authz_id->data, authz_id->size);

  return result;
}


add_if_prefix_matches(lookup_state_t *state,
                      const sorted_pattern_t *prefix,
                      const svn_stringbuf_t *segment)
{
  node_t *node = prefix->node;
  if (   node->segment.len <= segment->len
      && !memcmp(node->segment.data, segment->data, node->segment.len))
    add_next_node(state, node);
}


ensure_node_in_array(apr_array_header_t **array,
                     authz_rule_segment_t *segment,
                     apr_pool_t *result_pool)
{
  int idx;
  sorted_pattern_t entry;
  sorted_pattern_t *entry_ptr;

  /* Auto-create the array. */
  if (!*array)
    *array = apr_array_make(result_pool, 4, sizeof(sorted_pattern_t));

  /* Find the node in ARRAY and the IDX at which it were to be inserted.
   * Initialize IDX such that we won't attempt a hinted lookup (likely
   * to fail and therefore pure overhead). */
  idx = (*array)->nelts;
  entry_ptr = svn_sort__array_lookup(*array, segment, &idx,
                                     compare_node_rule_segment);
  if (entry_ptr)
    return entry_ptr->node;

  /* There is no such node, yet.
   * Create one and insert it into the sorted array. */
  entry.node = create_node(segment, result_pool);
  entry.next = NULL;
  svn_error_clear(svn_sort__array_insert2(*array, &entry, idx));

  return entry.node;
}


authz_read(authz_full_t **authz_p,
           svn_membuf_t **authz_id,
           const char *path,
           const char *groups_path,
           svn_boolean_t must_exist,
           svn_repos_t *repos_hint,
           svn_repos_authz_warning_func_t warning_func,
           void *warning_baton,
           apr_pool_t *result_pool,
           apr_pool_t *scratch_pool)
{
  svn_error_t* err = NULL;
  svn_stream_t *rules_stream = NULL;
  svn_stream_t *groups_stream = NULL;
  svn_checksum_t *rules_checksum = NULL;
  svn_checksum_t *groups_checksum = NULL;

  config_access_t *config_access =
    svn_repos__create_config_access(repos_hint, scratch_pool);

  /* Open the main authz file */
  SVN_ERR(svn_repos__get_config(&rules_stream, &rules_checksum, config_access,
                                path, must_exist, scratch_pool));

  /* Open the optional groups file */
  if (groups_path)
    SVN_ERR(svn_repos__get_config(&groups_stream, &groups_checksum,
                                  config_access, groups_path, must_exist,
                                  scratch_pool));

  /* The authz cache is optional. */
  *authz_id = construct_authz_key(rules_checksum, groups_checksum,
                                  result_pool);
  if (authz_pool)
    {
      /* Cache lookup. */
      SVN_ERR(svn_object_pool__lookup((void **)authz_p, authz_pool,
                                      *authz_id, result_pool));

      /* If not found, parse and add to cache. */
      if (!*authz_p)
        {
          apr_pool_t *item_pool = svn_object_pool__new_item_pool(authz_pool);

          /* Parse the configuration(s) and construct the full authz model
           * from it. */
          err = svn_authz__parse(authz_p, rules_stream, groups_stream,
                                 warning_func, warning_baton,
                                 item_pool, scratch_pool);
          if (err != SVN_NO_ERROR)
            {
              /* That pool would otherwise never get destroyed. */
              svn_pool_destroy(item_pool);

              /* Add the URL / file name to the error stack since the parser
               * doesn't have it. */
              err = svn_error_quick_wrapf(err,
                                   "Error while parsing config file: '%s':",
                                   path);
            }
          else
            {
              SVN_ERR(svn_object_pool__insert((void **)authz_p, authz_pool,
                                              *authz_id, *authz_p,
                                              item_pool, result_pool));
            }
        }
    }
  else
    {
      /* Parse the configuration(s) and construct the full authz model from
       * it. */
      err = svn_error_quick_wrapf(
          svn_authz__parse(authz_p, rules_stream, groups_stream,
                           warning_func, warning_baton,
                           result_pool, scratch_pool),
          "Error while parsing authz file: '%s':", path);
    }

  svn_repos__destroy_config_access(config_access);

  return err;
}


insert_path(construction_context_t *ctx,
            node_t *node,
            path_access_t *path_access,
            int segment_count,
            authz_rule_segment_t *segment,
            apr_pool_t *result_pool,
            apr_pool_t *scratch_pool)
{
  node_t *sub_node;
  node_segment_pair_t *node_segment;

  /* End of path? */
  if (segment_count == 0)
    {
      /* Set access rights.  Note that there might be multiple rules for
       * the same path due to non-repo-specific rules vs. repo-specific
       * ones.  Whichever gets defined last wins.
       */
      limited_rights_t rights;
      rights.access = *path_access;
      rights.max_rights = path_access->rights;
      rights.min_rights = path_access->rights;
      combine_access(&node->rights, &rights);
      return;
    }

  /* Any wildcards?  They will go into a separate sub-structure. */
  if (segment->kind != authz_rule_literal)
    ensure_pattern_sub_nodes(node, result_pool);

  switch (segment->kind)
    {
      /* A full wildcard segment? */
    case authz_rule_any_segment:
      sub_node = ensure_node(&node->pattern_sub_nodes->any,
                             segment, result_pool);
      break;

      /* One or more full wildcard segments? */
    case authz_rule_any_recursive:
      sub_node = ensure_node(&node->pattern_sub_nodes->any_var,
                             segment, result_pool);
      ensure_pattern_sub_nodes(sub_node, result_pool)->repeat = TRUE;
      break;

      /* A single wildcard at the end of the segment? */
    case authz_rule_prefix:
      sub_node = ensure_node_in_array(&node->pattern_sub_nodes->prefixes,
                                      segment, result_pool);
      break;

      /* A single wildcard at the start of segments? */
    case authz_rule_suffix:
      sub_node = ensure_node_in_array(&node->pattern_sub_nodes->suffixes,
                                      segment, result_pool);
      break;

      /* General pattern? */
    case authz_rule_fnmatch:
      sub_node = ensure_node_in_array(&node->pattern_sub_nodes->complex,
                                      segment, result_pool);
      break;

      /* Then it must be a literal. */
    default:
      SVN_ERR_ASSERT_NO_RETURN(segment->kind == authz_rule_literal);

      if (!node->sub_nodes)
        {
          node->sub_nodes = svn_hash__make(result_pool);
          sub_node = NULL;
        }
      else
        {
          sub_node = svn_hash_gets(node->sub_nodes, segment->pattern.data);
        }

      /* Auto-insert a sub-node for the current segment. */
      if (!sub_node)
        {
          sub_node = create_node(segment, result_pool);
          apr_hash_set(node->sub_nodes,
                       sub_node->segment.data,
                       sub_node->segment.len,
                       sub_node);
        }
    }

  /* Update context. */
  node_segment = apr_array_push(ctx->path);
  node_segment->segment = segment;
  node_segment->node = sub_node;

  /* Continue at the sub-node with the next segment. */
  insert_path(ctx, sub_node, path_access, segment_count - 1, segment + 1,
              result_pool, scratch_pool);
}


trim_subnode_hash(apr_hash_t **hash,
                  int latest_any_var,
                  apr_pool_t *scratch_pool)
{
  if (*hash)
    {
      apr_array_header_t *to_remove = apr_array_make(scratch_pool, 0,
                                                     sizeof(node_t *));

      apr_hash_index_t *hi;
      for (hi = apr_hash_first(scratch_pool, *hash);
           hi;
           hi = apr_hash_next(hi))
        {
          node_t *node = apr_hash_this_val(hi);
          if (trim_tree(node, latest_any_var, scratch_pool))
            APR_ARRAY_PUSH(to_remove, node_t *) = node;
        }

      /* Are some nodes left? */
      if (to_remove->nelts < apr_hash_count(*hash))
        {
          /* Remove empty nodes (if any). */
          int i;
          for (i = 0; i < to_remove->nelts; ++i)
            {
              node_t *node = APR_ARRAY_IDX(to_remove, i, node_t *);
              apr_hash_set(*hash, node->segment.data, node->segment.len,
                           NULL);
            }

          return FALSE;
        }

      /* No nodes left.  A NULL hash is more efficient than an empty one. */
      *hash = NULL;
    }

  return TRUE;
}


trim_subnode_array(apr_array_header_t **array,
                   int latest_any_var,
                   apr_pool_t *scratch_pool)
{
  if (*array)
    {
      int i, dest;
      for (i = 0, dest = 0; i < (*array)->nelts; ++i)
        {
          node_t *node = APR_ARRAY_IDX(*array, i, sorted_pattern_t).node;
          if (!trim_tree(node, latest_any_var, scratch_pool))
            {
              if (i != dest)
                APR_ARRAY_IDX(*array, dest, sorted_pattern_t)
                  = APR_ARRAY_IDX(*array, i, sorted_pattern_t);
              ++dest;
            }
        }

      /* Are some nodes left? */
      if (dest)
        {
          /* Trim it to the number of valid entries. */
          (*array)->nelts = dest;
          return FALSE;
        }

      /* No nodes left.  A NULL array is more efficient than an empty one. */
      *array = NULL;
    }

  return TRUE;
}


svn_repos_authz_parse2(svn_authz_t **authz_p,
                       svn_stream_t *stream,
                       svn_stream_t *groups_stream,
                       svn_repos_authz_warning_func_t warning_func,
                       void *warning_baton,
                       apr_pool_t *result_pool,
                       apr_pool_t *scratch_pool)
{
  svn_authz_t *authz = apr_pcalloc(result_pool, sizeof(*authz));
  authz->pool = result_pool;

  /* Parse the configuration and construct the full authz model from it. */
  SVN_ERR(svn_authz__parse(&authz->full, stream, groups_stream,
                           warning_func, warning_baton,
                           result_pool, scratch_pool));

  *authz_p = authz;
  return SVN_NO_ERROR;
}


init_lockup_state(lookup_state_t *state,
                  node_t *root,
                  const char *path)
{
  apr_size_t len = strlen(path);
  if (   (len > state->parent_path->len)
      && state->parent_path->len
      && (path[state->parent_path->len] == '/')
      && !memcmp(path, state->parent_path->data, state->parent_path->len))
    {
      /* The PARENT_PATH of the previous lookup is actually a parent path
       * of PATH.  The CURRENT node list already matches the parent path
       * and we only have to set the correct rights info. */
      state->rights = state->parent_rights;

      /* Tell the caller where to proceed. */
      return path + state->parent_path->len;
    }

  /* Start lookup at ROOT for the full PATH. */
  state->rights = root->rights;
  state->parent_rights = root->rights;

  apr_array_clear(state->next);
  apr_array_clear(state->current);
  APR_ARRAY_PUSH(state->current, node_t *) = root;

  /* Var-segment rules match empty segments as well */
  if (root->pattern_sub_nodes && root->pattern_sub_nodes->any_var)
   {
      node_t *node = root->pattern_sub_nodes->any_var;

      /* This is non-recursive due to ACL normalization. */
      combine_access(&state->rights, &node->rights);
      combine_right_limits(&state->rights, &node->rights);
      APR_ARRAY_PUSH(state->current, node_t *) = node;
   }

  svn_stringbuf_setempty(state->parent_path);
  svn_stringbuf_setempty(state->scratch_pad);

  return path;
}


create_lookup_state(apr_pool_t *result_pool)
{
  lookup_state_t *state = apr_pcalloc(result_pool, sizeof(*state));

  state->next = apr_array_make(result_pool, 4, sizeof(node_t *));
  state->current = apr_array_make(result_pool, 4, sizeof(node_t *));

  /* Virtually all path segments should fit into this buffer.  If they
   * don't, the buffer gets automatically reallocated.
   *
   * Using a smaller initial size would be fine as well but does not
   * buy us much for the increased risk of being expanded anyway - at
   * some extra cost. */
  state->scratch_pad = svn_stringbuf_create_ensure(200, result_pool);

  /* Most paths should fit into this buffer.  The same rationale as
   * above applies. */
  state->parent_path = svn_stringbuf_create_ensure(200, result_pool);

  return state;
}


combine_access(limited_rights_t *target,
               const limited_rights_t *rights)
{
  /* This implies the check for NO_SEQUENCE_NUMBER, i.e no rights being
   * specified. */
  if (target->access.sequence_number < rights->access.sequence_number)
    target->access = rights->access;
}


svn_repos_authz_initialize(apr_pool_t *pool)
{
  /* Protect against multiple calls. */
  return svn_error_trace(svn_atomic__init_once(&authz_pool_initialized,
                                               synchronized_authz_initialize,
                                               NULL, pool));
}


finalize_tree(node_t *node,
              limited_rights_t *sum,
              apr_pool_t *scratch_pool)
{
  limited_rights_t *local_sum = &node->rights;

  /* For convenience, we allow NODE to be NULL: */
  if (!node)
    return;

  /* Sum of rights at NODE - so far. */
  if (has_local_rule(local_sum))
    {
      local_sum->max_rights = local_sum->access.rights;
      local_sum->min_rights = local_sum->access.rights;
    }
  else
    {
      local_sum->min_rights = authz_access_write;
      local_sum->max_rights = authz_access_none;
    }

  /* Process all sub-nodes. */
  finalize_subnode_hash(node->sub_nodes, local_sum, scratch_pool);

  if (node->pattern_sub_nodes)
    {
      finalize_tree(node->pattern_sub_nodes->any, local_sum, scratch_pool);
      finalize_tree(node->pattern_sub_nodes->any_var, local_sum, scratch_pool);

      finalize_subnode_array(node->pattern_sub_nodes->prefixes, local_sum,
                             scratch_pool);
      finalize_subnode_array(node->pattern_sub_nodes->suffixes, local_sum,
                             scratch_pool);
      finalize_subnode_array(node->pattern_sub_nodes->complex, local_sum,
                             scratch_pool);

      /* Link up the prefixes / suffixes. */
      link_prefix_patterns(node->pattern_sub_nodes->prefixes);
      link_prefix_patterns(node->pattern_sub_nodes->suffixes);
    }

  /* Add our min / max info to the parent's info.
   * Idempotent for parent == node (happens at root). */
  combine_right_limits(sum, local_sum);
}


svn_repos_authz_read4(svn_authz_t **authz_p,
                      const char *path,
                      const char *groups_path,
                      svn_boolean_t must_exist,
                      svn_repos_t *repos_hint,
                      svn_repos_authz_warning_func_t warning_func,
                      void *warning_baton,
                      apr_pool_t *result_pool,
                      apr_pool_t *scratch_pool)
{
  svn_authz_t *authz = apr_pcalloc(result_pool, sizeof(*authz));
  authz->pool = result_pool;

  SVN_ERR(authz_read(&authz->full, &authz->authz_id, path, groups_path,
                     must_exist, repos_hint, warning_func, warning_baton,
                     result_pool, scratch_pool));

  *authz_p = authz;
  return SVN_NO_ERROR;
}


deinit_authz(void *data)
{
  /* The two object pools run their own cleanup handlers. */
  authz_pool = NULL;
  filtered_pool = NULL;
  authz_pool_initialized = FALSE;
  return APR_SUCCESS;
}


process_acl(construction_context_t *ctx,
            const authz_acl_t *acl,
            node_t *root,
            const char *repository,
            const char *user,
            apr_pool_t *result_pool,
            apr_pool_t *scratch_pool)
{
  path_access_t path_access;
  int i;
  node_t *node;

  /* Skip ACLs that don't say anything about the current user
     and/or repository. */
  if (!svn_authz__get_acl_access(&path_access.rights, acl, user, repository))
    return;

  /* Insert the rule into the filtered tree. */
  path_access.sequence_number = acl->sequence_number;

  /* Try to reuse results from previous runs.
   * Basically, skip the common prefix. */
  node = root;
  for (i = 0; i < ctx->path->nelts; ++i)
    {
      const node_segment_pair_t *step
        = &APR_ARRAY_IDX(ctx->path, i, const node_segment_pair_t);

      /* Exploit the fact that all strings in the authz model are unique /
       * internized and can be identified by address alone. */
      if (   !step->node
          || i >= acl->rule.len
          || step->segment->kind != acl->rule.path[i].kind
          || step->segment->pattern.data != acl->rule.path[i].pattern.data)
        {
          ctx->path->nelts = i;
          break;
        }
      else
        {
          node = step->node;
        }
    }

  /* Insert the path rule into the filtered tree. */
  insert_path(ctx, node, &path_access,
              acl->rule.len - i, acl->rule.path + i,
              result_pool, scratch_pool);
}


finalize_subnode_array(apr_array_header_t *array,
                       limited_rights_t *sum,
                       apr_pool_t *scratch_pool)
{
  if (array)
    {
      int i;
      for (i = 0; i < array->nelts; ++i)
        finalize_tree(APR_ARRAY_IDX(array, i, sorted_pattern_t).node, sum,
                      scratch_pool);
    }
}


filter_tree(svn_authz_t *authz,
            apr_pool_t *scratch_pool)
{
  apr_pool_t *pool = authz->filtered->pool;
  const char *repos_name = authz->filtered->repository;
  const char *user = authz->filtered->user;
  node_t *root;

  if (filtered_pool)
    {
      svn_membuf_t *key = construct_filtered_key(repos_name, user,
                                                 authz->authz_id,
                                                 scratch_pool);

      /* Cache lookup. */
      SVN_ERR(svn_object_pool__lookup((void **)&root, filtered_pool, key,
                                      pool));

      if (!root)
        {
          apr_pool_t *item_pool = svn_object_pool__new_item_pool(authz_pool);
          authz_full_t *add_ref = NULL;

          /* Make sure the underlying full authz object lives as long as the
           * filtered one that we are about to create.  We do this by adding
           * a reference to it in ITEM_POOL (which may live longer than AUTHZ).
           *
           * Note that we already have a reference to that full authz in
           * AUTHZ->FULL. Assert that we actually don't created multiple
           * instances of the same full model.
           */
          svn_error_clear(svn_object_pool__lookup((void **)&add_ref,
                                                  authz_pool, authz->authz_id,
                                                  item_pool));
          SVN_ERR_ASSERT(add_ref == authz->full);

          /* Now construct the new filtered tree and cache it. */
          root = create_user_authz(authz->full, repos_name, user, item_pool,
                                   scratch_pool);
          svn_error_clear(svn_object_pool__insert((void **)&root,
                                                  filtered_pool, key, root,
                                                  item_pool, pool));
        }
     }
  else
    {
      root = create_user_authz(authz->full, repos_name, user, pool,
                               scratch_pool);
    }

  /* Write a new entry. */
  authz->filtered->root = root;

  return SVN_NO_ERROR;
}


construct_authz_key(const svn_checksum_t *authz_key,
                    const svn_checksum_t *groups_key,
                    apr_pool_t *result_pool)
{
  svn_membuf_t *result = apr_pcalloc(result_pool, sizeof(*result));
  if (groups_key)
    {
      apr_size_t authz_size = svn_checksum_size(authz_key);
      apr_size_t groups_size = svn_checksum_size(groups_key);

      svn_membuf__create(result, authz_size + groups_size, result_pool);
      result->size = authz_size + groups_size; /* exact length is required! */

      memcpy(result->data, authz_key->digest, authz_size);
      memcpy((char *)result->data + authz_size,
             groups_key->digest, groups_size);
    }
  else
    {
      apr_size_t size = svn_checksum_size(authz_key);
      svn_membuf__create(result, size, result_pool);
      result->size = size; /* exact length is required! */
      memcpy(result->data, authz_key->digest, size);
    }

  return result;
}


create_construction_context(apr_pool_t *result_pool)
{
  construction_context_t *result = apr_pcalloc(result_pool, sizeof(*result));

  /* Array will be auto-extended but this initial size will make it rarely
   * ever necessary. */
  result->path = apr_array_make(result_pool, 32, sizeof(node_segment_pair_t));

  return result;
}


compare_node_rule_segment(const void *void_lhs,
                          const void *void_rhs)
{
  const sorted_pattern_t *element = void_lhs;
  const authz_rule_segment_t *segment = void_rhs;

  return strcmp(element->node->segment.data, segment->pattern.data);
}


trim_tree(node_t *node,
          int latest_any_var,
          apr_pool_t *scratch_pool)
{
  svn_boolean_t removed_all = TRUE;

  /* For convenience, we allow NODE to be NULL: */
  if (!node)
    return TRUE;

  /* Do we have a later "any_var" rule at this node. */
  if (   node->pattern_sub_nodes
      && node->pattern_sub_nodes->any_var
      &&   node->pattern_sub_nodes->any_var->rights.access.sequence_number
         > latest_any_var)
    {
      latest_any_var
        = node->pattern_sub_nodes->any_var->rights.access.sequence_number;
    }

  /* Is there a local rule at this node that is not eclipsed by any_var? */
  if (has_local_rule(&node->rights))
    {
      /* Remove the local rule, if it got eclipsed.
       * Note that for the latest any_var node, the sequence number is equal. */
      if (node->rights.access.sequence_number >= latest_any_var)
        removed_all = FALSE;
      else
         node->rights.access.sequence_number = NO_SEQUENCE_NUMBER;
    }

  /* Process all sub-nodes. */
  removed_all &= trim_subnode_hash(&node->sub_nodes, latest_any_var,
                                   scratch_pool);

  if (node->pattern_sub_nodes)
    {
      if (trim_tree(node->pattern_sub_nodes->any, latest_any_var,
                    scratch_pool))
        node->pattern_sub_nodes->any = NULL;
      else
        removed_all = FALSE;

      if (trim_tree(node->pattern_sub_nodes->any_var, latest_any_var,
                    scratch_pool))
        node->pattern_sub_nodes->any_var = NULL;
      else
        removed_all = FALSE;

      removed_all &= trim_subnode_array(&node->pattern_sub_nodes->prefixes,
                                        latest_any_var, scratch_pool);
      removed_all &= trim_subnode_array(&node->pattern_sub_nodes->suffixes,
                                        latest_any_var, scratch_pool);
      removed_all &= trim_subnode_array(&node->pattern_sub_nodes->complex,
                                        latest_any_var, scratch_pool);

      /* Trim the tree as much as possible to speed up lookup(). */
      if (removed_all)
        node->pattern_sub_nodes = NULL;
    }

  return removed_all;
}


synchronized_authz_initialize(void *baton, apr_pool_t *pool)
{
#if APR_HAS_THREADS
  svn_boolean_t multi_threaded = TRUE;
#else
  svn_boolean_t multi_threaded = FALSE;
#endif

  SVN_ERR(svn_object_pool__create(&authz_pool, multi_threaded, pool));
  SVN_ERR(svn_object_pool__create(&filtered_pool, multi_threaded, pool));

  apr_pool_cleanup_register(pool, NULL, deinit_authz, apr_pool_cleanup_null);
  return SVN_NO_ERROR;
}


svn_repos_authz_check_access(svn_authz_t *authz, const char *repos_name,
                             const char *path, const char *user,
                             svn_repos_authz_access_t required_access,
                             svn_boolean_t *access_granted,
                             apr_pool_t *pool)
{
  const authz_access_t required =
    ((required_access & svn_authz_read ? authz_access_read_flag : 0)
     | (required_access & svn_authz_write ? authz_access_write_flag : 0));

  /* Pick or create the suitable pre-filtered path rule tree. */
  authz_user_rules_t *rules = get_user_rules(
      authz,
      (repos_name ? repos_name : AUTHZ_ANY_REPOSITORY),
      user);

  /* In many scenarios, users have uniform access to a repository
   * (blanket access or no access at all).
   *
   * In these cases, don't bother creating or consulting the filtered tree.
   */
  if ((rules->global_rights.min_access & required) == required)
    {
      *access_granted = TRUE;
      return SVN_NO_ERROR;
    }

  if ((rules->global_rights.max_access & required) != required)
    {
      *access_granted = FALSE;
      return SVN_NO_ERROR;
    }

  /* No specific path given, i.e. looking for anywhere in the tree? */
  if (!path)
    {
      *access_granted =
        ((rules->global_rights.max_access & required) == required);
      return SVN_NO_ERROR;
    }

  /* Rules tree lookup */

  /* Did we already filter the data model? */
  if (!rules->root)
    SVN_ERR(filter_tree(authz, pool));

  /* Re-use previous lookup results, if possible. */
  path = init_lockup_state(authz->filtered->lookup_state,
                           authz->filtered->root, path);

  /* Sanity check. */
  SVN_ERR_ASSERT(path[0] == '/');

  /* Determine the granted access for the requested path.
   * PATH does not need to be normalized for lockup(). */
  *access_granted = lookup(rules->lookup_state, path, required,
                           !!(required_access & svn_authz_recursive), pool);

  return SVN_NO_ERROR;
}


create_user_authz(authz_full_t *authz,
                  const char *repository,
                  const char *user,
                  apr_pool_t *result_pool,
                  apr_pool_t *scratch_pool)
{
  int i;
  node_t *root = create_node(NULL, result_pool);
  construction_context_t *ctx = create_construction_context(scratch_pool);

  /* Use a separate sub-pool to keep memory usage tight. */
  apr_pool_t *subpool = svn_pool_create(scratch_pool);

  /* Find all ACLs for REPOSITORY. */
  apr_array_header_t *acls = apr_array_make(subpool, authz->acls->nelts,
                                            sizeof(authz_acl_t *));
  for (i = 0; i < authz->acls->nelts; ++i)
    {
      const authz_acl_t *acl = &APR_ARRAY_IDX(authz->acls, i, authz_acl_t);
      if (svn_authz__acl_applies_to_repo(acl, repository))
        {
          /* ACLs in the AUTHZ are sorted by path and repository.
           * So, if there is a rule for the repo and a global rule for the
           * same path, we will detect them here. */
          if (acls->nelts)
            {
              const authz_acl_t *prev_acl
                = APR_ARRAY_IDX(acls, acls->nelts - 1, const authz_acl_t *);
              if (svn_authz__compare_paths(&prev_acl->rule, &acl->rule) == 0)
                {
                  svn_boolean_t global_acl_applies;
                  svn_boolean_t repos_acl_applies;

                  /* Previous ACL is a global rule. */
                  SVN_ERR_ASSERT_NO_RETURN(!strcmp(prev_acl->rule.repos,
                                                   AUTHZ_ANY_REPOSITORY));
                  /* Current ACL is a per-repository rule. */
                  SVN_ERR_ASSERT_NO_RETURN(strcmp(acl->rule.repos,
                                                  AUTHZ_ANY_REPOSITORY));

                  global_acl_applies =
                    svn_authz__get_acl_access(NULL, prev_acl, user, repository);
                  repos_acl_applies =
                    svn_authz__get_acl_access(NULL, acl, user, repository);

                  /* Prefer rules which apply to both this user and this path
                   * over rules which apply only to the path. In cases where
                   * both rules apply to user and path, always prefer the
                   * repository-specific rule. */
                  if (!global_acl_applies || repos_acl_applies)
                    {
                      apr_array_pop(acls);
                      APR_ARRAY_PUSH(acls, const authz_acl_t *) = acl;
                    }
                }
              else
                APR_ARRAY_PUSH(acls, const authz_acl_t *) = acl;
            }
          else
            APR_ARRAY_PUSH(acls, const authz_acl_t *) = acl;
        }
    }

  /* Filtering and tree construction. */
  for (i = 0; i < acls->nelts; ++i)
    process_acl(ctx, APR_ARRAY_IDX(acls, i, const authz_acl_t *),
                root, repository, user, result_pool, subpool);

  /* If there is no relevant rule at the root node, the "no access" default
   * applies. Give it a SEQUENCE_NUMBER that will never overrule others. */
  if (!has_local_rule(&root->rights))
    {
      root->rights.access.sequence_number = 0;
      root->rights.access.rights = authz_access_none;
    }

  /* Trim the tree.
   *
   * We can't do pattern comparison, so for most pattern rules we cannot
   * say that a set of rules "eclipses" / overrides a given other set of
   * rules for all possible paths.  That limits the accuracy of our check
   * for recursive access in similar ways than for non-pattern rules.
   *
   * However, the user expects a rule ending with "**" to eclipse any older
   * rule in that sub-tree recursively.  So, this trim function removes
   * eclipsed nodes from the tree.
   */
  svn_pool_clear(subpool);
  trim_tree(root, NO_SEQUENCE_NUMBER, subpool);

  /* Calculate recursive rights.
   *
   * This is a bottom-up calculation of the range of access rights
   * specified anywhere in  the respective sub-tree, including the base
   * node itself.
   *
   * To prevent additional finalization passes, we piggy-back the addition
   * of the ordering links of the prefix and suffix sub-node rules.
   */
  svn_pool_clear(subpool);
  finalize_tree(root, &root->rights, subpool);

  /* Done. */
  svn_pool_destroy(subpool);
  return root;
}


add_prefix_matches(lookup_state_t *state,
                   const svn_stringbuf_t *segment,
                   apr_array_header_t *prefixes)
{
  /* Index of the first node that might be a match.  All matches will
   * be at this and the immediately following indexes. */
  int i = svn_sort__bsearch_lower_bound(prefixes, segment->data,
                                        compare_node_path_segment);

  /* The entry we found may be an exact match (but not a true prefix).
   * The prefix matching test will still work. */
  if (i < prefixes->nelts)
    add_if_prefix_matches(state,
                          &APR_ARRAY_IDX(prefixes, i, sorted_pattern_t),
                          segment);

  /* The immediate predecessor may be a true prefix and all potential
   * prefixes can be found following the NEXT links between the array
   * indexes. */
  if (i > 0)
    {
      sorted_pattern_t *pattern;
      for (pattern = &APR_ARRAY_IDX(prefixes, i - 1, sorted_pattern_t);
           pattern;
           pattern = pattern->next)
        {
          add_if_prefix_matches(state, pattern, segment);
        }
    }
}


combine_right_limits(limited_rights_t *target,
                     const limited_rights_t *rights)
{
  target->max_rights |= rights->max_rights;
  target->min_rights &= rights->min_rights;
}


  void Compute(OpKernelContext* ctx) override {
    try {
      const Tensor& input = ctx->input(kInputTensorIndex);
      OP_REQUIRES(
          ctx, input.dims() == 4,
          errors::InvalidArgument("Current RequantizePerChannel operator"
                                  "supports 4D tensors only."));

      const Tensor& input_min_vec = ctx->input(kInputMinVecIndex);
      size_t depth = input_min_vec.NumElements();
      float* input_min_vec_data = (float*)const_cast<void*>(
          static_cast<const void*>(input_min_vec.flat<float>().data()));

      const Tensor& input_max_vec = ctx->input(kInputMaxVecIndex);
      OP_REQUIRES(
          ctx, input_max_vec.NumElements() == depth,
          errors::InvalidArgument("input_max has incorrect size, expected ",
                                  depth, " was ", input_max_vec.NumElements()));
      float* input_max_vec_data = (float*)const_cast<void*>(
          static_cast<const void*>(input_max_vec.flat<float>().data()));

      const Tensor& input_requested_min = ctx->input(this->kRequestMinIndex);
      OP_REQUIRES(
          ctx, input_requested_min.NumElements() == 1,
          errors::InvalidArgument("requested_output_min must be a scalar"));
      const float input_requested_min_float =
          input_requested_min.flat<float>()(0);

      const Tensor& input_requested_max = ctx->input(this->kRequestMaxIndex);
      OP_REQUIRES(
          ctx, input_requested_min.NumElements() == 1,
          errors::InvalidArgument("requested_output_max must be a scalar"));
      const float input_requested_max_float =
          input_requested_max.flat<float>()(0);

      if (out_type_ == DT_QINT8) {
        OP_REQUIRES(ctx, input_requested_min_float < 0.0f,
                    errors::InvalidArgument(
                        "If out_type is QINT8, requested_output_max must be "
                        "non negative, got ",
                        input_requested_min_float));
      }

      const float factor = (out_type_ == DT_QINT8) ? 127.0f : 255.0f;
      const float requested_min_max =
          std::max(std::abs(input_requested_min_float),
                   std::abs(input_requested_max_float));
      Tensor* output = nullptr;
      OP_REQUIRES_OK(ctx, ctx->allocate_output(kOutputTensorIndex,
                                               input.shape(), &output));

      std::vector<float> scales(depth);
      for (int i = 0; i < depth; ++i) {
        float min_max_from_vec = std::max(std::abs(input_min_vec_data[i]),
                                          std::abs(input_max_vec_data[i]));
        scales[i] = factor * (min_max_from_vec / requested_min_max /
                              static_cast<float>(1L << 31));
      }

      mkldnn::primitive_attr reorder_attr;
      reorder_attr.set_output_scales(2, scales);

      memory::dims dims_mkl_order =
          TFShapeToMklDnnDimsInNCHW(input.shape(), FORMAT_NHWC);
      memory::desc input_md = memory::desc(dims_mkl_order, MklDnnType<qint32>(),
                                           memory::format_tag::nhwc);
      memory::desc output_md =
          (out_type_ == DT_QINT8)
              ? memory::desc(dims_mkl_order, MklDnnType<qint8>(),
                             memory::format_tag::nhwc)
              : memory::desc(dims_mkl_order, MklDnnType<quint8>(),
                             memory::format_tag::nhwc);

      void* input_buf =
          static_cast<void*>(const_cast<qint32*>(input.flat<qint32>().data()));
      void* output_buf;
      if (out_type_ == DT_QINT8) {
        output_buf = static_cast<void*>(
            const_cast<qint8*>(output->flat<qint8>().data()));
      } else {
        output_buf = static_cast<void*>(
            const_cast<quint8*>(output->flat<quint8>().data()));
      }

      std::unique_ptr<memory> input_mem_prim(
          new memory(input_md, cpu_engine_, input_buf));
      std::unique_ptr<memory> output_mem_prim(
          new memory(output_md, cpu_engine_, output_buf));

      mkldnn::reorder::primitive_desc reorder_pd =
          ReorderPd(cpu_engine_, input_mem_prim->get_desc(), cpu_engine_,
                    output_mem_prim->get_desc(), reorder_attr);
      std::shared_ptr<stream> reorder_stream;
      MklDnnThreadPool eigen_tp(ctx);
      reorder_stream.reset(CreateStream(&eigen_tp, cpu_engine_));
      std::unordered_map<int, mkldnn::memory> reorder_args = {
          {MKLDNN_ARG_FROM, *input_mem_prim},
          {MKLDNN_ARG_TO, *output_mem_prim}};
      std::unique_ptr<mkldnn::primitive> reorder_prim(
          new mkldnn::reorder(reorder_pd));
      reorder_prim->execute(*reorder_stream, reorder_args);

      Tensor* output_min = nullptr;
      Tensor* output_max = nullptr;
      OP_REQUIRES_OK(ctx,
                     ctx->allocate_output(kOutputMinIndex, {}, &output_min));
      OP_REQUIRES_OK(ctx,
                     ctx->allocate_output(kOutputMaxIndex, {}, &output_max));

      output_min->flat<float>()(0) = input_requested_min_float;
      output_max->flat<float>()(0) = input_requested_max_float;
    } catch (mkldnn::error& e) {
      string error_msg = "Status: " + std::to_string(e.status) +
                         ", message: " + std::string(e.message) + ", in file " +
                         std::string(__FILE__) + ":" + std::to_string(__LINE__);
      OP_REQUIRES_OK(
          ctx, errors::Aborted("Operation received an exception:", error_msg));
    }
  }

Reorganize and add more validation to MKL requantization

PiperOrigin-RevId: 387901341
Change-Id: I2515b9034c64e113db0bcec8337d30643ab0a0f1


  explicit MklRequantizePerChannelOp(OpKernelConstruction* ctx)
      : OpKernel(ctx) {
    OP_REQUIRES_OK(ctx, ctx->GetAttr("out_type", &out_type_));
    OP_REQUIRES(ctx, out_type_ == DT_QINT8 || out_type_ == DT_QUINT8,
                errors::InvalidArgument(
                    "out_type must be qint8 or quint8, but got: ", out_type_));
  }


TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
  auto* params = reinterpret_cast<TfLiteSVDFParams*>(node->builtin_data);
  OpData* op_data = reinterpret_cast<OpData*>(node->user_data);

  const TfLiteTensor* input;
  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kInputTensor, &input));
  const TfLiteTensor* weights_feature;
  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kWeightsFeatureTensor,
                                          &weights_feature));
  const TfLiteTensor* weights_time;
  TF_LITE_ENSURE_OK(
      context, GetInputSafe(context, node, kWeightsTimeTensor, &weights_time));
  const TfLiteTensor* bias = GetOptionalInputTensor(context, node, kBiasTensor);

  TfLiteTensor* scratch;
  TF_LITE_ENSURE_OK(context,
                    GetTemporarySafe(context, node, /*index=*/0, &scratch));

  TfLiteTensor* state = GetVariableInput(context, node, kStateTensor);
  TF_LITE_ENSURE(context, state != nullptr);
  TfLiteTensor* output;
  TF_LITE_ENSURE_OK(context,
                    GetOutputSafe(context, node, kOutputTensor, &output));

  switch (weights_feature->type) {
    case kTfLiteFloat32: {
      reference_ops::EvalFloatSVDF(
          params, GetTensorShape(input), GetTensorData<float>(input),
          GetTensorShape(weights_feature),
          GetTensorData<float>(weights_feature), GetTensorShape(weights_time),
          GetTensorData<float>(weights_time), GetTensorShape(bias),
          GetTensorData<float>(bias), GetTensorData<float>(scratch),
          GetTensorData<float>(state), GetTensorShape(output),
          GetTensorData<float>(output));
      return kTfLiteOk;
    }
    case kTfLiteUInt8:
    case kTfLiteInt8: {
      if (input->type == kTfLiteFloat32) {
        TfLiteTensor* input_quantized;
        TF_LITE_ENSURE_OK(context, GetTemporarySafe(context, node, /*index=*/1,
                                                    &input_quantized));
        TfLiteTensor* scaling_factors;
        TF_LITE_ENSURE_OK(context, GetTemporarySafe(context, node, /*index=*/2,
                                                    &scaling_factors));
        TfLiteTensor* float_weights_time;
        TF_LITE_ENSURE_OK(context, GetTemporarySafe(context, node, /*index=*/3,
                                                    &float_weights_time));
        TfLiteTensor* zero_points;
        TF_LITE_ENSURE_OK(context, GetTemporarySafe(context, node, /*index=*/4,
                                                    &zero_points));
        TfLiteTensor* row_sums;
        TF_LITE_ENSURE_OK(
            context, GetTemporarySafe(context, node, /*index=*/5, &row_sums));
        // Dequantize weights time.
        // TODO(alanchiao): this dequantization initialization only needs to
        // happen once per model and should theoretically be placed in either
        // Init or Prepare. However, TFLite doesn't allocate float_weights_time
        // until the Eval function.
        // TODO(alanchiao): refactor logic out into dequantize function.
        if (!op_data->float_weights_time_initialized) {
          const float dequantization_scale = weights_time->params.scale;
          const int8_t* weights_time_ptr = GetTensorData<int8_t>(weights_time);
          float* float_weights_time_ptr =
              GetTensorData<float>(float_weights_time);
          for (int i = 0; i < NumElements(float_weights_time); ++i) {
            float_weights_time_ptr[i] =
                weights_time_ptr[i] * dequantization_scale;
          }
          op_data->float_weights_time_initialized = true;
        }

        int32_t* zero_points_ptr = nullptr;
        int32_t* row_sums_ptr = nullptr;
        if (params->asymmetric_quantize_inputs && row_sums != nullptr) {
          zero_points_ptr = GetTensorData<int32_t>(zero_points);
          row_sums_ptr = GetTensorData<int32_t>(row_sums);
        }

        reference_ops::EvalHybridSVDF(
            params, GetTensorShape(input), GetTensorData<float>(input),
            GetTensorShape(weights_feature),
            GetTensorData<int8_t>(weights_feature),
            weights_feature->params.scale, GetTensorShape(float_weights_time),
            GetTensorData<float>(float_weights_time), GetTensorShape(bias),
            GetTensorData<float>(bias), GetTensorData<float>(scratch),
            GetTensorData<float>(scaling_factors),
            GetTensorData<int8_t>(input_quantized), GetTensorData<float>(state),
            GetTensorShape(output), GetTensorData<float>(output),
            zero_points_ptr, row_sums_ptr, &op_data->compute_row_sums);
        return kTfLiteOk;
      }
      auto* input_params = reinterpret_cast<TfLiteAffineQuantization*>(
          input->quantization.params);
      auto* output_params = reinterpret_cast<TfLiteAffineQuantization*>(
          output->quantization.params);
      TfLiteTensor* output_temp;
      TF_LITE_ENSURE_OK(
          context, GetTemporarySafe(context, node, /*index=*/1, &output_temp));

      // Currently supports only ReLU.
      // TODO(jianlijianli): support other activations.
      TF_LITE_ENSURE_EQ(context, params->activation, kTfLiteActRelu);

      reference_ops::EvalIntegerSVDF(
          params, GetTensorShape(input), GetTensorData<int8_t>(input),
          GetTensorShape(weights_feature),
          GetTensorData<int8_t>(weights_feature), GetTensorShape(weights_time),
          GetTensorData<int16_t>(weights_time), GetTensorShape(bias),
          GetTensorData<int32_t>(bias), GetTensorData<int16_t>(state),
          GetTensorShape(output), GetTensorData<int8_t>(output),
          GetTensorData<int32_t>(scratch), GetTensorData<int32_t>(output_temp),
          op_data->effective_scale_1_a, op_data->effective_scale_1_b,
          op_data->effective_scale_2_a, op_data->effective_scale_2_b,
          input_params->zero_point->data[0],
          output_params->zero_point->data[0]);
      return kTfLiteOk;
    }
    default:
      context->ReportError(context, "Type %s not currently supported.",
                           TfLiteTypeGetName(weights_feature->type));
      return kTfLiteError;
  }
}

Fix a null pointer exception caused by branching on uninitialized data.

This is due to not checking that the params for the quantization exists. If there is no quantization, we should not access the `.params` field.

PiperOrigin-RevId: 385163909
Change-Id: I2beb8d50649b6542db224c163033fbcbaa49314f


TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
  const auto* params = reinterpret_cast<TfLiteSVDFParams*>(node->builtin_data);
  OpData* op_data = reinterpret_cast<OpData*>(node->user_data);
  int scratch_tensor_index = op_data->scratch_tensor_index;

  // Check we have all the inputs and outputs we need.
  TF_LITE_ENSURE_EQ(context, node->outputs->size, 1);
  TF_LITE_ENSURE_EQ(context, node->inputs->size, 5);

  const TfLiteTensor* input;
  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kInputTensor, &input));
  const TfLiteTensor* weights_feature;
  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kWeightsFeatureTensor,
                                          &weights_feature));
  const TfLiteTensor* weights_time;
  TF_LITE_ENSURE_OK(
      context, GetInputSafe(context, node, kWeightsTimeTensor, &weights_time));

  TF_LITE_ENSURE(context,
                 input->type == kTfLiteFloat32 || input->type == kTfLiteInt8);

  // Check all the parameters of tensor match within themselves and match the
  // input configuration.
  const int rank = params->rank;
  const int batch_size = input->dims->data[0];
  const int num_filters = weights_feature->dims->data[0];
  TF_LITE_ENSURE(context, rank != 0);
  TF_LITE_ENSURE_EQ(context, num_filters % rank, 0);
  const int num_units = num_filters / rank;
  const int memory_size = weights_time->dims->data[1];
  TF_LITE_ENSURE_EQ(context, input->dims->data[1],
                    weights_feature->dims->data[1]);
  TF_LITE_ENSURE_EQ(context, weights_time->dims->data[0], num_filters);

  const TfLiteTensor* bias = GetOptionalInputTensor(context, node, kBiasTensor);
  if (bias) {
    TF_LITE_ENSURE_EQ(context, bias->dims->data[0], num_units);
  }

  const TfLiteTensor* state;
  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kStateTensor, &state));
  TfLiteTensor* output;
  TF_LITE_ENSURE_OK(context,
                    GetOutputSafe(context, node, kOutputTensor, &output));

  // Check the shape of input state tensors.
  TF_LITE_ENSURE_EQ(context, NumDimensions(state), 2);
  TF_LITE_ENSURE_EQ(context, SizeOfDimension(state, 0), batch_size);
  TF_LITE_ENSURE_EQ(context, SizeOfDimension(state, 1),
                    memory_size * num_filters);

  // Resize output.
  TfLiteIntArray* output_size_array = TfLiteIntArrayCreate(2);
  output_size_array->data[0] = batch_size;
  output_size_array->data[1] = num_units;
  TF_LITE_ENSURE_OK(context,
                    context->ResizeTensor(context, output, output_size_array));

  // The weights are of consistent type, so it suffices to check one.
  const bool is_hybrid_op = IsHybridOp(input, weights_feature);
  const bool is_full_integer = input->type == kTfLiteInt8;

  // Resize scratch.
  TfLiteIntArrayFree(node->temporaries);
  if (is_hybrid_op) {
    node->temporaries = TfLiteIntArrayCreate(6);
  } else if (is_full_integer) {
    node->temporaries = TfLiteIntArrayCreate(2);
  } else {
    node->temporaries = TfLiteIntArrayCreate(1);
  }
  node->temporaries->data[0] = scratch_tensor_index;

  TfLiteIntArray* scratch_size_array = TfLiteIntArrayCreate(2);
  scratch_size_array->data[0] = batch_size;
  scratch_size_array->data[1] = num_filters;

  TfLiteTensor* scratch_tensor;
  TF_LITE_ENSURE_OK(
      context, GetTemporarySafe(context, node, /*index=*/0, &scratch_tensor));

  // The scratch buffer is of type int32 for full integer svdf and it's of type
  // float32 for hybrid and float case.
  if (is_full_integer) {
    scratch_tensor->type = kTfLiteInt32;
  } else {
    scratch_tensor->type = kTfLiteFloat32;
  }
  scratch_tensor->allocation_type = kTfLiteArenaRw;
  TF_LITE_ENSURE_OK(context, context->ResizeTensor(context, scratch_tensor,
                                                   scratch_size_array));

  if (is_hybrid_op) {
    op_data->compute_row_sums = true;
    // Tell interpreter to allocate temporary tensors to store quantized values
    // of input tensors.
    node->temporaries->data[1] = scratch_tensor_index + 1;
    TfLiteTensor* input_quantized;
    TF_LITE_ENSURE_OK(context, GetTemporarySafe(context, node, /*index=*/1,
                                                &input_quantized));
    input_quantized->type = weights_feature->type;
    input_quantized->allocation_type = kTfLiteArenaRw;
    if (!TfLiteIntArrayEqual(input_quantized->dims, input->dims)) {
      TfLiteIntArray* input_quantized_size = TfLiteIntArrayCopy(input->dims);
      TF_LITE_ENSURE_OK(context, context->ResizeTensor(context, input_quantized,
                                                       input_quantized_size));
    }

    // Tell interpreter to allocate temporary tensors to store scaling factors.
    node->temporaries->data[2] = scratch_tensor_index + 2;
    TfLiteTensor* scaling_factors;
    TF_LITE_ENSURE_OK(context, GetTemporarySafe(context, node, /*index=*/2,
                                                &scaling_factors));
    scaling_factors->type = kTfLiteFloat32;
    scaling_factors->allocation_type = kTfLiteArenaRw;
    int scaling_dims[1] = {batch_size};
    if (!TfLiteIntArrayEqualsArray(scaling_factors->dims, 1, scaling_dims)) {
      TfLiteIntArray* scaling_factors_size = TfLiteIntArrayCreate(1);
      scaling_factors_size->data[0] = batch_size;
      TF_LITE_ENSURE_OK(context, context->ResizeTensor(context, scaling_factors,
                                                       scaling_factors_size));
    }

    // Used to store dequantized weights_time matrix for hybrid computation of
    // matmul(state, weights_time), which occurs in floating point.
    node->temporaries->data[3] = scratch_tensor_index + 3;
    TfLiteTensor* float_weights_time;
    TF_LITE_ENSURE_OK(context, GetTemporarySafe(context, node, /*index=*/3,
                                                &float_weights_time));
    float_weights_time->type = kTfLiteFloat32;
    // Persistent so that we can compute the dequantized weights only once.
    float_weights_time->allocation_type = kTfLiteArenaRwPersistent;
    if (!TfLiteIntArrayEqual(float_weights_time->dims, weights_time->dims)) {
      TfLiteIntArray* float_weights_time_size =
          TfLiteIntArrayCopy(weights_time->dims);
      TF_LITE_ENSURE_OK(context,
                        context->ResizeTensor(context, float_weights_time,
                                              float_weights_time_size));
    }

    node->temporaries->data[4] = scratch_tensor_index + 4;
    TfLiteTensor* zero_points;
    TF_LITE_ENSURE_OK(
        context, GetTemporarySafe(context, node, /*index=*/4, &zero_points));
    zero_points->type = kTfLiteFloat32;
    zero_points->allocation_type = kTfLiteArenaRw;
    int zero_points_dims[1] = {batch_size};
    if (!TfLiteIntArrayEqualsArray(zero_points->dims, 1, zero_points_dims)) {
      TfLiteIntArray* zero_points_size = TfLiteIntArrayCreate(1);
      zero_points_size->data[0] = zero_points_dims[0];
      TF_LITE_ENSURE_OK(context, context->ResizeTensor(context, zero_points,
                                                       zero_points_size));
    }

    node->temporaries->data[5] = scratch_tensor_index + 5;
    TfLiteTensor* row_sums;
    TF_LITE_ENSURE_OK(context,
                      GetTemporarySafe(context, node, /*index=*/5, &row_sums));
    row_sums->type = kTfLiteFloat32;
    row_sums->allocation_type = kTfLiteArenaRwPersistent;
    int row_sums_dims[1] = {num_filters};
    if (!TfLiteIntArrayEqualsArray(row_sums->dims, 1, row_sums_dims)) {
      TfLiteIntArray* row_sums_size = TfLiteIntArrayCreate(1);
      row_sums_size->data[0] = row_sums_dims[0];
      TF_LITE_ENSURE_OK(
          context, context->ResizeTensor(context, row_sums, row_sums_size));
    }
  }
  if (is_full_integer) {
    // Allocated one extra tensor.
    TfLiteIntArray* output_temp_size_array = TfLiteIntArrayCreate(2);
    output_temp_size_array->data[0] = num_units;
    output_temp_size_array->data[1] = batch_size;
    node->temporaries->data[1] = scratch_tensor_index + 1;
    TfLiteTensor* output_temp;
    TF_LITE_ENSURE_OK(
        context, GetTemporarySafe(context, node, /*index=*/1, &output_temp));
    output_temp->type = kTfLiteInt32;
    output_temp->allocation_type = kTfLiteArenaRw;
    TF_LITE_ENSURE_OK(context, context->ResizeTensor(context, output_temp,
                                                     output_temp_size_array));

    // Calculate effective scales.
    TF_LITE_ENSURE(context, input->quantization.type != kTfLiteNoQuantization);
    auto* input_params =
        reinterpret_cast<TfLiteAffineQuantization*>(input->quantization.params);
    TF_LITE_ENSURE(context,
                   weights_feature->quantization.type != kTfLiteNoQuantization);
    auto* weights_feature_params = reinterpret_cast<TfLiteAffineQuantization*>(
        weights_feature->quantization.params);
    TF_LITE_ENSURE(context, state->quantization.type != kTfLiteNoQuantization);
    auto* state_params =
        reinterpret_cast<TfLiteAffineQuantization*>(state->quantization.params);
    TF_LITE_ENSURE(context,
                   weights_time->quantization.type != kTfLiteNoQuantization);
    auto* weight_time_params = reinterpret_cast<TfLiteAffineQuantization*>(
        weights_time->quantization.params);
    TF_LITE_ENSURE(context, output->quantization.type != kTfLiteNoQuantization);
    auto* output_params = reinterpret_cast<TfLiteAffineQuantization*>(
        output->quantization.params);
    const double effective_scale_1 = input_params->scale->data[0] *
                                     weights_feature_params->scale->data[0] /
                                     state_params->scale->data[0];
    const double effective_scale_2 = state_params->scale->data[0] *
                                     weight_time_params->scale->data[0] /
                                     output_params->scale->data[0];
    QuantizeMultiplier(effective_scale_1, &op_data->effective_scale_1_a,
                       &op_data->effective_scale_1_b);
    QuantizeMultiplier(effective_scale_2, &op_data->effective_scale_2_a,
                       &op_data->effective_scale_2_b);
  }
  return kTfLiteOk;
}


Status ShapeRefiner::SetShape(const Node* node, int output_port,
                              ShapeHandle shape) {
  auto c = GetContext(node);
  if (c == nullptr) {
    return errors::Internal("Could not find context for ", node->name());
  }

  if (output_port < 0 || output_port >= node->num_outputs()) {
    return errors::InvalidArgument(
        "output_port '", output_port, "' is out of range, ", "node '",
        node->name(), "' has ", node->num_outputs(), " outputs");
  }
  // Note: it's possible, if the node's been updated, that the shape inference
  // context doesn't have the right number of outputs.
  if (node->num_outputs() > c->num_outputs()) {
    TF_RETURN_IF_ERROR(c->ExpandOutputs(node->num_outputs()));
  }

  // Check compatibility, and merge the shapes.
  ShapeHandle existing_shape = c->output(output_port);
  TF_RETURN_IF_ERROR(c->Merge(existing_shape, shape, &shape));
  c->set_output(output_port, shape);

  // TODO(vrv): Do we need to propagate the new shape through all
  // consumers that change their outputs?  At the moment, python
  // does not do this, but this seems like a nice feature.

  // TODO(vrv): We might need to keep track of the fact that the
  // existing shape is invalidated, in case we need to propagate
  // this information to remote workers.
  return Status::OK();
}

Fix segmentation fault in shape inference logic.

When running shape functions, some functions (such as `MutableHashTableShape`)
produce extra output information in the form of a `ShapeAndType` struct.  The
shapes embedded in this struct are owned by an inference context that is
cleaned up almost immediately; if the upstream code attempts to access this
shape information, it can trigger a segfault.

`ShapeRefiner` is mitigating this for normal output shapes by cloning them
(and thus putting the newly created shape under ownership of an inference
context that will not die), but we were not doing the same for shapes and
types.  This commit fixes that by doing similar logic on output shapes and
types.

PiperOrigin-RevId: 384761124
Change-Id: I07c0c42d29dfbb55bfa13ec1f09ef825fb0a1a1d


Status ShapeRefiner::RunShapeFn(const Node* node,
                                const OpRegistrationData* op_reg_data,
                                ExtendedInferenceContext* ec,
                                InferenceContext* outer_context) {
  // This will be filled in with real data in a second pass.
  std::vector<const Tensor*> input_tensors(node->num_inputs(), nullptr);
  std::vector<Tensor> real_tensors(node->num_inputs());
  std::vector<bool> attempted_materialization(node->num_inputs());
  std::vector<bool> attempted_tensor_as_shape_conversion(node->num_inputs());
  std::vector<ShapeHandle> input_tensors_as_shapes;

  auto* c = ec->get_context();

  c->set_input_tensors(input_tensors);
  c->set_input_tensors_as_shapes(input_tensors_as_shapes);

  // Run the shape inference function, and return if there was an error.
  // Capture as lambda, because we might need to re-run inference later on.
  auto run_inference_lambda = [&]() {
    if (function_library_ && IsFunctionCall(*function_library_, *node)) {
      bool disable_shape_inference;
      if (!GetNodeAttr(AttrSlice(node->def()), "_disable_call_shape_inference",
                       &disable_shape_inference)
               .ok() ||
          !disable_shape_inference) {
        // Special inference logic for user-defined functions.
        NameAttrList function;
        TF_RETURN_IF_ERROR(
            NameAndAttrsFromFunctionCall(node->def(), &function));
        const FunctionDef* function_def =
            function_library_->Find(function.name());
        if (function_def != nullptr) {
          // The constant Tensor map we have for the outside context is not
          // valid inside the function. We need to push a new clean map while
          // performing inference on the function body.
          auto const_tensor_map_copy = const_tensor_map_;
          const_tensor_map_.clear();
          Status function_inference_status = InferShapesForFunction(
              function_def, AttrSlice(&function.attr()), ec);
          const_tensor_map_ = const_tensor_map_copy;
          return function_inference_status;
        }
      }
    }

    if (op_reg_data->shape_inference_fn) {
      TF_RETURN_IF_ERROR(c->Run(op_reg_data->shape_inference_fn));
    } else {
      TF_RETURN_IF_ERROR(c->Run(shape_inference::UnknownShape));
    }
    return Status::OK();
  };
  TF_RETURN_IF_ERROR(run_inference_lambda());

  // We must run the shape function repeatedly, in case users write
  // shape functions where they only conditionally call input_tensor()
  // based on the values of another input tensor.
  bool rerun_shape_fn;
  do {
    // If the result of running shape inference would have benefitted
    // from knowing the values of input tensors, try to materialize
    // the results of those tensors, and then run the shape inference
    // function again using those known tensors.
    rerun_shape_fn = false;

    // NOTE: It is possible to batch the extraction and
    // materialization of inputs, instead of materializing one input
    // at a time like we do below.  If input-at-a-time computation
    // becomes a bottleneck, we could separate ExtractConstantSubgraph
    // into two functions: one that returns true if an input is
    // derivable from constants, and another function that extracts
    // the subgraph for multiple target nodes and executes the whole
    // subgraph once.

    for (int i = 0; i < c->num_inputs(); ++i) {
      if (!c->requested_input_tensor(i)) {
        continue;
      }
      // Check if we have not already filled in the requested input,
      // and if not, try to materialize the tensors.
      if (!attempted_materialization[i]) {
        attempted_materialization[i] = true;

        Tensor result;
        bool evaluated = false;
        TF_RETURN_IF_ERROR(EvaluateConstantTensorForEdge(
            node, i, &evaluated, &result, outer_context));
        if (evaluated) {
          real_tensors[i] = result;
          input_tensors[i] = &real_tensors[i];
          // We have more concrete information about a shape,
          // so re-run shape inference.
          rerun_shape_fn = true;
        }
      }
      if (c->requested_input_tensor_as_partial_shape(i) &&
          !attempted_tensor_as_shape_conversion[i]) {
        attempted_tensor_as_shape_conversion[i] = true;
        if (i >= input_tensors_as_shapes.size()) {
          input_tensors_as_shapes.resize(i + 1);
        }
        ShapeHandle s;
        TF_RETURN_IF_ERROR(ConstantPartialShape(c, node, i, &s, outer_context));
        input_tensors_as_shapes[i] = s;
        rerun_shape_fn = true;
      }
    }

    if (rerun_shape_fn) {
      // We have more information about the shapes on this pass,
      // so re-run shape inference.
      c->set_input_tensors(input_tensors);
      c->set_input_tensors_as_shapes(input_tensors_as_shapes);
      TF_RETURN_IF_ERROR(run_inference_lambda());
    }
  } while (rerun_shape_fn);

  return Status::OK();
}


ShapeRefiner::~ShapeRefiner() {
  // The lifetime of the tensors are bound to the GraphRunner, so the tensors
  // should be deleted before it.
  const_tensor_map_.clear();
}


ShapeRefiner::ShapeRefiner(const VersionDef& versions,
                           const OpRegistryInterface* ops)
    : ShapeRefiner(versions.producer(), ops) {}


Status ShapeRefiner::PartialStridedSliceShape(
    Node* slice_node, InferenceContext* ctx, ShapeHandle* result,
    shape_inference::InferenceContext* outer_context) {
  // Only attempt to evaluate if begin/end/strides all are scalars.
  for (int i = 1; i <= 3; ++i) {
    ShapeHandle input_shape = ctx->input(i);
    if (ctx->Value(ctx->Dim(input_shape, 0)) != 1) {
      *result = ctx->UnknownShape();
      return Status::OK();
    }
  }

  int begin_mask, end_mask, ellipsis_mask, new_axis_mask, shrink_axis_mask;
  TF_RETURN_IF_ERROR(
      GetNodeAttr(slice_node->attrs(), "begin_mask", &begin_mask));
  TF_RETURN_IF_ERROR(GetNodeAttr(slice_node->attrs(), "end_mask", &end_mask));
  TF_RETURN_IF_ERROR(
      GetNodeAttr(slice_node->attrs(), "ellipsis_mask", &ellipsis_mask));
  TF_RETURN_IF_ERROR(
      GetNodeAttr(slice_node->attrs(), "new_axis_mask", &new_axis_mask));
  TF_RETURN_IF_ERROR(
      GetNodeAttr(slice_node->attrs(), "shrink_axis_mask", &shrink_axis_mask));

  // Only attempt to evaluate if there are no special masks set (note that we
  // can handle begin/end_mask == 1).
  if (!(begin_mask == 0 || begin_mask == 1) ||
      !(end_mask == 0 || end_mask == 1) || ellipsis_mask != 0 ||
      new_axis_mask != 0 || shrink_axis_mask != 0) {
    *result = ctx->UnknownShape();
    return Status::OK();
  }

  bool evaluated;
  int64 begin;
  if (begin_mask == 1) {
    begin = 0;
  } else {
    TF_RETURN_IF_ERROR(EvaluateConstantIntScalarEdge(slice_node, 1, &evaluated,
                                                     &begin, outer_context));
    if (!evaluated) {
      *result = ctx->UnknownShape();
      return Status::OK();
    }
  }

  int64 end;
  if (end_mask == 1) {
    end = std::numeric_limits<int64>::max();
  } else {
    TF_RETURN_IF_ERROR(EvaluateConstantIntScalarEdge(slice_node, 2, &evaluated,
                                                     &end, outer_context));
    if (!evaluated) {
      *result = ctx->UnknownShape();
      return Status::OK();
    }
  }

  int64 stride;
  TF_RETURN_IF_ERROR(EvaluateConstantIntScalarEdge(slice_node, 3, &evaluated,
                                                   &stride, outer_context));
  if (!evaluated) {
    *result = ctx->UnknownShape();
    return Status::OK();
  }

  // Apply stride to input interpreted as a partial shape.
  ShapeHandle input;
  TF_RETURN_IF_ERROR(
      ConstantPartialShape(ctx, slice_node, 0, &input, outer_context));
  TF_RETURN_IF_ERROR(ctx->Subshape(input, begin, end, stride, result));
  return Status::OK();
}


bool ShapeRefiner::IsUpdatedShapesOrTypes(
    InferenceContext* c, const std::vector<ShapeAndType>& existing,
    const std::vector<ShapeAndType>& updated) {
  if (existing.size() != updated.size()) {
    return true;
  }
  for (int i = 0; i < existing.size(); i++) {
    if (!SameDefinedShape(c, existing[i].shape, updated[i].shape) ||
        existing[i].dtype != updated[i].dtype) {
      return true;
    }
  }
  return false;
}


Status ShapeRefiner::InferShapesForFunctionSubNode(
    const Node* node, InferenceContext* outer_context) {
  TF_RETURN_IF_ERROR(AddNodeInternal(node, outer_context));
  InferenceContext* node_context = CHECK_NOTNULL(GetContext(node));

  if (StringPiece(node->type_string()) == kArgOp) {
    // Handle special node: function input.
    // Shapes for these nodes are provided in the outer inference
    // context.

    int index;
    TF_RETURN_IF_ERROR(GetNodeAttr(AttrSlice(node->def()), "index", &index));

    if (index < 0 || outer_context->num_inputs() <= index) {
      return errors::Internal(
          "Function instantiation included invalid input index: ", index,
          " not in [0, ", outer_context->num_inputs(), ").");
    }

    // TODO(b/134547156): TEMPORARY WORKAROUND. If input shape handle is not set
    // in outer context, set _Arg node output shape to unknown.
    if (outer_context->input(index).SameHandle(ShapeHandle())) {
      VLOG(1) << "Function instantiation has undefined input shape at "
              << "index: " << index << " in the outer inference context.";
      node_context->set_output(0, node_context->UnknownShape());
    } else {
      node_context->set_output(0, outer_context->input(index));
    }

    auto* resource = outer_context->input_handle_shapes_and_types(index);
    if (resource) {
      node_context->set_output_handle_shapes_and_types(0, *resource);
    }
  } else if (StringPiece(node->type_string()) == kRetvalOp) {
    // Handle special node: function output.
    // Shapes inferred for these nodes go into the outer inference
    // context.

    int index;
    TF_RETURN_IF_ERROR(GetNodeAttr(AttrSlice(node->def()), "index", &index));

    if (index < 0 || outer_context->num_outputs() <= index) {
      return errors::Internal(
          "Function instantiation included invalid output index: ", index,
          " not in [0, ", outer_context->num_outputs(), ").");
    }

    // outer_context outlives node_context, therefore we need to create
    // a new shape handle owned by outer_context instead.
    ShapeHandle handle;
    TensorShapeProto proto;
    node_context->ShapeHandleToProto(node_context->input(0), &proto);
    TF_RETURN_IF_ERROR(outer_context->MakeShapeFromShapeProto(proto, &handle));
    outer_context->set_output(index, handle);

    const std::vector<ShapeAndType>* resource =
        node_context->input_handle_shapes_and_types(0);
    if (resource) {
      // `ShapesAndType`s contain `ShapeHandle`s.  These `ShapeHandle`s point
      // to `Shape`s that are owned by a different inference context too.  We
      // need to copy them to the outer context to prevent them from being
      // destroyed before they are used.
      std::vector<ShapeAndType> copied_shapes_and_types;
      for (auto& shape_and_type : *resource) {
        ShapeHandle handle;
        TensorShapeProto proto;
        node_context->ShapeHandleToProto(shape_and_type.shape, &proto);
        TF_RETURN_IF_ERROR(
            outer_context->MakeShapeFromShapeProto(proto, &handle));
        copied_shapes_and_types.push_back(
            ShapeAndType(handle, shape_and_type.dtype, shape_and_type.type));
      }

      outer_context->set_output_handle_shapes_and_types(
          index, copied_shapes_and_types);
    }
  }

  return Status::OK();
}


Status ShapeRefiner::UpdateNode(const Node* node, bool relax, bool* refined) {
  auto it = node_to_context_.find(node);
  if (it == node_to_context_.end()) {
    *refined = true;
    return AddNode(node);
  }
  ExtendedInferenceContext* node_ext_context = it->second.get();
  InferenceContext* node_context = node_ext_context->get_context();

  // Give up if the context wasn't successfully built by the AddNode() method.
  TF_RETURN_IF_ERROR(node_context->construction_status());

  // Check if the shapes of the nodes in the fan-in of this node have changed,
  // and if they have update the node input shapes.
  for (const Edge* e : node->in_edges()) {
    if (e->IsControlEdge()) continue;

    int dst_input = e->dst_input();
    int src_output = e->src_output();

    Node* input = e->src();
    auto iter = node_to_context_.find(input);
    if (iter == node_to_context_.end()) {
      return errors::FailedPrecondition(
          "Input ", dst_input, " ('", input->name(), "') for '", node->name(),
          "' was not previously added to ShapeRefiner.");
    }

    InferenceContext* c = iter->second->get_context();
    DCHECK_GE(dst_input, 0);
    ShapeHandle existing_input = node_context->input(dst_input);
    if (!relax) {
      if (node_context->MergeInput(dst_input, c->output(src_output))) {
        if (!SameDefinedShape(node_context, node_context->input(dst_input),
                              existing_input)) {
          *refined = true;
        }
      }
    } else {
      if (node_context->RelaxInput(dst_input, c->output(src_output))) {
        if (!SameDefinedShape(node_context, node_context->input(dst_input),
                              existing_input)) {
          *refined = true;
        }
      }
    }
    if (node_context->requested_input_tensor_as_partial_shape(dst_input)) {
      // The input value may have changed. Since we have no way to know if
      // that's indeed the case, err on the safe side.
      *refined = true;
    }

    // Also propagate handle shape and dtype of edges which are carrying
    // resource handles.
    if (e->src()->output_type(src_output) == DT_RESOURCE) {
      auto* outputs = c->output_handle_shapes_and_types(src_output);
      if (!outputs) continue;

      if (!relax &&
          node_context->MergeInputHandleShapesAndTypes(dst_input, *outputs)) {
        *refined = true;
      } else if (relax) {
        std::vector<ShapeAndType> existing_inputs;
        const std::vector<ShapeAndType>* inputs =
            node_context->input_handle_shapes_and_types(dst_input);
        if (inputs) {
          existing_inputs = *inputs;
        }
        if (node_context->RelaxInputHandleShapesAndMergeTypes(dst_input,
                                                              *outputs)) {
          if (IsUpdatedShapesOrTypes(
                  node_context, existing_inputs,
                  *node_context->input_handle_shapes_and_types(dst_input))) {
            *refined = true;
          }
        }
      }
    }
  }

  if (!*refined) {
    // No input shape has changed, we're done
    return Status::OK();
  }

  // Get and run the shape function for this node to update the shapes of the
  // outputs.
  const OpRegistrationData* op_reg_data;
  TF_RETURN_IF_ERROR(ops_registry_->LookUp(node->type_string(), &op_reg_data));
  if (op_reg_data->shape_inference_fn == nullptr &&
      require_shape_inference_fns_) {
    return errors::InvalidArgument(
        "No shape inference function exists for op '", node->type_string(),
        "', did you forget to define it?");
  }

  if (!op_reg_data->shape_inference_fn) {
    // There is nothing more we can infer
    return Status::OK();
  }

  return RunShapeFn(node, op_reg_data, node_ext_context);
}


Status ShapeRefiner::ConstantPartialShape(
    InferenceContext* target_context, const Node* node, int dst_idx,
    ShapeHandle* result, shape_inference::InferenceContext* outer_context) {
  const Edge* input_edge;
  TF_RETURN_IF_ERROR(node->input_edge(dst_idx, &input_edge));

  InferenceContext* src_context = GetContext(input_edge->src());
  if (src_context == nullptr) return errors::Internal("Missing src context");
  ShapeHandle src_shape = src_context->output(input_edge->src_output());

  if (src_context->Value(src_context->Rank(src_shape)) == 0) {
    Tensor t;
    bool evaluated = false;
    TF_RETURN_IF_ERROR(EvaluateConstantTensorForEdge(node, dst_idx, &evaluated,
                                                     &t, outer_context));
    if (!evaluated) {
      return errors::InvalidArgument(
          "Received a shape scalar with unknown static value.  A static value "
          "of '-1' is required to represent an unknown shape.");
    }
    if (t.dims() == 0) {
      if (t.dtype() == DT_INT32 && t.scalar<int32>()() == -1) {
        *result = target_context->UnknownShape();
        return Status::OK();
      } else if (t.dtype() == DT_INT64 && t.scalar<int64>()() == -1) {
        *result = target_context->UnknownShape();
        return Status::OK();
      }
    }
    return errors::InvalidArgument(
        "Received an invalid shape scalar with a static value that is not "
        "'-1': ",
        t.DebugString());
  }

  TF_RETURN_IF_ERROR(src_context->WithRank(src_shape, 1, &src_shape));

  const string& src_op = input_edge->src()->type_string();
  if (src_context->Value(src_context->Dim(src_shape, 0)) == 0) {
    // Source tensor is a vector of length 0, so the shape it
    // represents is as scalar.
    *result = target_context->Scalar();
  } else if (src_op == "Cast") {
    // First try to evaluate the current tensor, as it might be a valid cast of
    // a float.
    Tensor t;
    bool evaluated = false;
    if (EvaluateConstantTensorForEdge(node, dst_idx, &evaluated, &t,
                                      outer_context)
            .ok()) {
      if (evaluated &&
          target_context->MakeShapeFromTensor(&t, src_shape, result).ok()) {
        return Status::OK();
      }
    }

    // Then try to infer partial shape from the input to the cast tensor.
    ShapeHandle pre_cast_shape;
    if (!ConstantPartialShape(target_context, input_edge->src(), 0,
                              &pre_cast_shape, outer_context)
             .ok()) {
      TF_RETURN_IF_ERROR(
          target_context->MakeShapeFromTensor(nullptr, src_shape, result));
    }
    if (!target_context->RankKnown(pre_cast_shape)) {
      // Failed to evaluate. Treat the output as completely unknown.
      *result = target_context->UnknownShape();
      return Status::OK();
    }
    auto* dest_type = input_edge->src()->attrs().Find("DstT");
    if (dest_type == nullptr || dest_type->value_case() != AttrValue::kType ||
        (dest_type->type() != DT_INT32 && dest_type->type() != DT_INT64)) {
      // Casting to a weird type. Do not attempt to infer across it.
      *result = target_context->MakeShape(std::vector<DimensionHandle>(
          target_context->Rank(pre_cast_shape), target_context->UnknownDim()));
      return Status::OK();
    }
    *result = pre_cast_shape;
  } else if (src_op == "Shape") {
    *result = src_context->input(0);
  } else if (src_op == "ShapeN") {
    *result = src_context->input(input_edge->src_output());
  } else if (src_op == "Pack") {
    std::vector<DimensionHandle> dims;
    // Pack is concatenating its input scalars to form the shape tensor vector.
    for (int i = 0; i < src_context->num_inputs(); ++i) {
      int64 size;
      bool evaluated;
      TF_RETURN_IF_ERROR(EvaluateConstantIntScalarEdge(
          input_edge->src(), i, &evaluated, &size, outer_context));
      if (evaluated) {
        dims.push_back(size < 0 ? target_context->UnknownDim()
                                : target_context->MakeDim(size));
      } else {
        dims.push_back(target_context->UnknownDim());
      }
    }
    *result = target_context->MakeShape(dims);
  } else if (src_op == "Concat" || src_op == "ConcatV2") {
    *result = target_context->Scalar();
    // For Concat, input 0 is concat dim; for V2 it is the last input.
    const int concat_dim =
        src_op == "Concat" ? 0 : src_context->num_inputs() - 1;
    // Concat is concatenating its input shape vectors.
    for (int i = 0; i < src_context->num_inputs(); ++i) {
      // Concat dim is ignored (and will always be a scalar).
      if (i == concat_dim) continue;
      ShapeHandle sub_result;
      TF_RETURN_IF_ERROR(ConstantPartialShape(target_context, input_edge->src(),
                                              i, &sub_result, outer_context));
      if (!target_context->RankKnown(sub_result)) {
        // Failed to evaluate. Treat the output as completely unknown.
        // TODO(cwhipkey): we could rely on all inputs being the same rank, so
        // figure that rank out and append the right number of unknown dims.
        *result = target_context->UnknownShape();
        return Status::OK();
      }
      TF_RETURN_IF_ERROR(
          target_context->Concatenate(*result, sub_result, result));
    }
  } else if (src_op == "StridedSlice") {
    TF_RETURN_IF_ERROR(PartialStridedSliceShape(input_edge->src(), src_context,
                                                result, outer_context));
  } else if (src_op == "VariableShape") {
    auto* handle_data = src_context->input_handle_shapes_and_types(0);
    if (handle_data != nullptr && !handle_data->empty()) {
      *result = handle_data->at(0).shape;
    } else {
      *result = target_context->UnknownShape();
    }
  } else {
    Tensor t;
    bool evaluated = false;
    TF_RETURN_IF_ERROR(EvaluateConstantTensorForEdge(node, dst_idx, &evaluated,
                                                     &t, outer_context));
    TF_RETURN_IF_ERROR(target_context->MakeShapeFromTensor(
        evaluated ? &t : nullptr, src_shape, result));
  }
  return Status::OK();
}


Status ShapeRefiner::InferShapesForFunction(
    const FunctionDef* function_def, AttrSlice attributes,
    ExtendedInferenceContext* outer_context) {
  const Graph* graph;
  auto it = functions_.find(function_def);
  if (it != functions_.end()) {
    graph = it->second.get();
  } else {
    InstantiationResult result;
    TF_RETURN_IF_ERROR(InstantiateFunction(
        *function_def, attributes,
        [this](const string& op, const OpDef** sig) {
          return this->function_library_->LookUpOpDef(op, sig);
        },
        &result));

    Graph* new_graph = new Graph(function_library_);
    GraphConstructorOptions options;
    options.allow_internal_ops = true;
    TF_RETURN_IF_ERROR(
        ConvertNodeDefsToGraph(options, result.nodes, new_graph));
    functions_[function_def].reset(new_graph);
    graph = new_graph;
  }

  std::unordered_set<const Node*> function_nodes;
  Status inference_status = Status::OK();
  {
    auto node_shape_inference_lambda = [this, &outer_context, &function_nodes,
                                        &inference_status](const Node* node) {
      if (!inference_status.ok()) return;
      inference_status =
          InferShapesForFunctionSubNode(node, outer_context->get_context());
      function_nodes.insert(node);
    };

    // Calls inference lambda for each node after visiting all predecessors.
    // Ensures that we are adding nodes to ShapeRefiner in the topological
    // order.
    ReverseDFS(*graph, {}, node_shape_inference_lambda);
  }

  // Delete the contexts created for the functions nodes to save memory.
  for (const Node* node : function_nodes) {
    node_to_context_.erase(node);
  }

  return inference_status;
}


Status ShapeRefiner::EvaluateConstantTensorForEdge(
    const Node* node, int dst_idx, bool* evaluated, Tensor* result,
    InferenceContext* outer_context) {
  *evaluated = false;
  const Edge* input_edge;
  TF_RETURN_IF_ERROR(node->input_edge(dst_idx, &input_edge));
  OutputTensor tensor(input_edge->src(), input_edge->src_output());
  return EvaluateConstantTensor(
      tensor, *this, *ops_registry_, graph_def_version_, evaluated, result,
      &graph_runner_, &const_tensor_map_, kMaxTensorSize,
      disable_constant_propagation_, outer_context);
}


Status ShapeRefiner::AddNodeInternal(
    const Node* node, shape_inference::InferenceContext* outer_context) {
  // Create the inference context for this node with the existing input shapes.
  std::unique_ptr<InferenceContext> ic(new InferenceContext(
      graph_def_version_, node->def(), node->op_def(),
      std::vector<ShapeHandle>(node->num_inputs()), {}, {}, {}));
  TF_RETURN_IF_ERROR(ic->construction_status());

  // For each 'input' of this node, fetch the corresponding shape
  // from 'input's InferenceContext, and store into this node's
  // InferenceContext.
  for (const Edge* e : node->in_edges()) {
    if (e->IsControlEdge()) continue;

    if (e->dst_input() < 0) {
      return tensorflow::errors::Internal(
          "Index ", e->dst_input(), " is negative but not a control edge.");
    }

    const Node* input = e->src();
    auto it = node_to_context_.find(input);
    if (it == node_to_context_.end()) {
      // v1 control flow adds loops to the graph; we have to break them
      // somewhere, so we'll ignore this input and leave its shape undefined.
      ic->SetInput(e->dst_input(), ic->UnknownShape());
      continue;
    }

    InferenceContext* input_ic = it->second->get_context();
    ic->SetInput(e->dst_input(), input_ic->output(e->src_output()));

    const auto* in_v =
        input_ic->output_handle_shapes_and_types(e->src_output());
    if (in_v != nullptr) {
      DataType input_type = e->src()->output_type(e->src_output());
      DCHECK(input_type == DT_RESOURCE || input_type == DT_VARIANT);
      ic->set_input_handle_shapes_and_types(e->dst_input(),
                                            std::vector<ShapeAndType>(*in_v));
    }
  }

  // Get the shape function for this node
  const OpRegistrationData* op_reg_data;
  TF_RETURN_IF_ERROR(ops_registry_->LookUp(node->type_string(), &op_reg_data));
  if (op_reg_data->shape_inference_fn == nullptr &&
      require_shape_inference_fns_) {
    return errors::InvalidArgument(
        "No shape inference function exists for op '", node->type_string(),
        "', did you forget to define it?");
  }

  std::unique_ptr<ExtendedInferenceContext> ec(
      new ExtendedInferenceContext(std::move(ic), node));

  // Run the shape inference function, and return if there was an error.
  TF_RETURN_IF_ERROR(RunShapeFn(node, op_reg_data, ec.get(), outer_context));

  // Store the resulting context object in the map.
  node_to_context_[node].swap(ec);

  return Status::OK();
}


bool ShapeRefiner::SameDefinedShape(InferenceContext* c, ShapeHandle s0,
                                    ShapeHandle s1) {
  if (s0.SameHandle(s1)) {
    return true;
  }
  if (c->Rank(s0) != c->Rank(s1)) {
    return false;
  }
  if (!c->RankKnown(s0) && !c->RankKnown(s1)) {
    return false;
  }
  for (int i = 0; i < c->Rank(s0); ++i) {
    if (!c->Dim(s0, i).SameHandle(c->Dim(s1, i))) {
      int64 val0 = c->Value(c->Dim(s0, i));
      int64 val1 = c->Value(c->Dim(s1, i));
      if (val0 < 0 || val1 < 0 || val0 != val1) {
        return false;
      }
    }
  }

  return true;
}


ShapeRefiner::ShapeRefiner(int graph_def_version,
                           const OpRegistryInterface* ops)
    : graph_def_version_(graph_def_version),
      ops_registry_(ops),
      graph_runner_(Env::Default()) {}


Status ShapeRefiner::AddNode(const Node* node) {
  return AddNodeInternal(node, /*outer_context=*/nullptr);
}