distributed_point_function.cc | Explore in Territory

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

#include "dpf/distributed_point_function.h"

#include <algorithm>
#include <array>
#include <cstddef>
#include <limits>
#include <memory>
#include <numeric>
#include <string>
#include <utility>
#include <vector>

#include "absl/container/btree_map.h"
#include "absl/container/flat_hash_map.h"
#include "absl/container/inlined_vector.h"
#include "absl/log/absl_check.h"
#include "absl/log/absl_log.h"
#include "absl/memory/memory.h"
#include "absl/numeric/int128.h"
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/string_view.h"
#include "absl/types/span.h"
#include "dpf/internal/evaluate_prg_hwy.h"
#include "dpf/internal/get_hwy_mode.h"
#include "dpf/internal/proto_validator.h"
#include "dpf/internal/value_type_helpers.h"
#include "dpf/status_macros.h"
#include "google/protobuf/io/zero_copy_stream_impl_lite.h"
#include "hwy/aligned_allocator.h"
#include "openssl/rand.h"

namespace distributed_point_functions {

namespace {

// PRG keys used to expand seeds using AES. The first two are used to compute
// correction words of seeds, while the last is used to compute correction
// words of the incremental DPF values. Values were computed by taking the
// first half of the SHA256 sum of the constant name, e.g., `echo
// "DistributedPointFunction::kPrgKeyLeft" | sha256sum`
constexpr absl::uint128 kPrgKeyLeft = …;
constexpr absl::uint128 kPrgKeyRight = …;
constexpr absl::uint128 kPrgKeyValue = …;

}  // namespace

DistributedPointFunction::DistributedPointFunction(
    std::unique_ptr<dpf_internal::ProtoValidator> proto_validator,
    std::vector<int> blocks_needed, Aes128FixedKeyHash prg_left,
    Aes128FixedKeyHash prg_right, Aes128FixedKeyHash prg_value,
    absl::flat_hash_map<std::string, ValueCorrectionFunction>
        value_correction_functions)
    : … { … }

absl::StatusOr<std::vector<Value>>
DistributedPointFunction::ComputeValueCorrection(
    int hierarchy_level, absl::Span<const absl::uint128> seeds,
    absl::uint128 alpha, const Value& beta, bool invert) const { … }

// Expands the PRG seeds at the next `tree_level`, updates `seeds` and
// `control_bits`, and writes the next correction word to `keys`.
absl::Status DistributedPointFunction::GenerateNext(
    int tree_level, absl::uint128 alpha, absl::Span<const Value> beta,
    absl::Span<absl::uint128> seeds, absl::Span<bool> control_bits,
    absl::Span<DpfKey> keys) const { … }

absl::uint128 DistributedPointFunction::DomainToTreeIndex(
    absl::uint128 domain_index, int hierarchy_level) const { … }

int DistributedPointFunction::DomainToBlockIndex(absl::uint128 domain_index,
                                                 int hierarchy_level) const { … }

absl::Status DistributedPointFunction::EvaluateSeeds(
    absl::Span<const absl::uint128> seeds, absl::Span<const bool> control_bits,
    absl::Span<const absl::uint128> paths,
    absl::Span<const CorrectionWord* const> correction_words,
    absl::Span<absl::uint128> seeds_out,
    absl::Span<bool> control_bits_out) const { … }

absl::StatusOr<DistributedPointFunction::DpfExpansion>
DistributedPointFunction::ExpandSeeds(
    const DpfExpansion& partial_evaluations,
    absl::Span<const CorrectionWord* const> correction_words) const { … }

absl::StatusOr<DistributedPointFunction::DpfExpansion>
DistributedPointFunction::ComputePartialEvaluations(
    absl::Span<const absl::uint128> prefixes, int hierarchy_level,
    bool update_ctx, EvaluationContext& ctx) const { … }

absl::StatusOr<DistributedPointFunction::DpfExpansion>
DistributedPointFunction::ExpandAndUpdateContext(
    int hierarchy_level, absl::Span<const absl::uint128> prefixes,
    EvaluationContext& ctx) const { … }

absl::StatusOr<hwy::AlignedFreeUniquePtr<absl::uint128[]>>
DistributedPointFunction::HashExpandedSeeds(
    int hierarchy_level, absl::Span<const absl::uint128> expansion) const { … }

absl::StatusOr<std::string>
DistributedPointFunction::SerializeValueTypeDeterministically(
    const ValueType& value_type) { … }

absl::StatusOr<DistributedPointFunction::ValueCorrectionFunction>
DistributedPointFunction::GetValueCorrectionFunction(
    const DpfParameters& parameters) const { … }

absl::StatusOr<std::unique_ptr<DistributedPointFunction>>
DistributedPointFunction::Create(const DpfParameters& parameters) { … }

absl::StatusOr<std::unique_ptr<DistributedPointFunction>>
DistributedPointFunction::CreateIncremental(
    absl::Span<const DpfParameters> parameters) { … }

absl::StatusOr<std::pair<DpfKey, DpfKey>>
DistributedPointFunction::GenerateKeysIncremental(
    absl::uint128 alpha, absl::Span<const Value> beta) { … }

absl::StatusOr<EvaluationContext>
DistributedPointFunction::CreateEvaluationContext(DpfKey key) const { … }

}  // namespace distributed_point_functions
chromium/third_party/distributed_point_functions/code/dpf/distributed_point_function.cc
