<commit_before>/*
  Copyright (c) 2018, 2019 Jouni Siren

  Author: Jouni Siren <jouni.siren@iki.fi>

  Permission is hereby granted, free of charge, to any person obtaining a copy
  of this software and associated documentation files (the "Software"), to deal
  in the Software without restriction, including without limitation the rights
  to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
  copies of the Software, and to permit persons to whom the Software is
  furnished to do so, subject to the following conditions:

  The above copyright notice and this permission notice shall be included in all
  copies or substantial portions of the Software.

  THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
  IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
  FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
  AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
  LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
  OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
  SOFTWARE.
*/

#include <string>
#include <unistd.h>

#include <gbwt/dynamic_gbwt.h>

using namespace gbwt;

//------------------------------------------------------------------------------

const std::string tool_name = "GBWT sequence remove";

void printUsage(int exit_code = EXIT_SUCCESS);

//------------------------------------------------------------------------------

int
main(int argc, char** argv)
{
  if(argc < 3) { printUsage(); }

  int c = 0;
  std::string output;
  size_type chunk_size = DynamicGBWT::REMOVE_CHUNK_SIZE;
  bool range = false;
  while((c = getopt(argc, argv, "c:o:r")) != -1)
  {
    switch(c)
    {
    case 'c':
      chunk_size = std::max(1ul, std::stoul(optarg));
      break;
    case 'o':
      output = optarg; break;
    case 'r':
      range = true; break;
    case '?':
      std::exit(EXIT_FAILURE);
    default:
      std::exit(EXIT_FAILURE);
    }
  }

  if(optind + 1 >= argc) { printUsage(EXIT_FAILURE); }
  std::string base_name = argv[optind]; optind++;
  if(output.empty()) { output = base_name; }
  std::vector<size_type> seq_ids;
  if(range)
  {
    if(argc != optind + 2) { printUsage(EXIT_FAILURE); }
    size_type start = std::stoul(argv[optind]); optind++;
    size_type stop = std::stoul(argv[optind]); optind++;
    if(stop < start) { printUsage(EXIT_FAILURE); }
    for(size_type seq_id = start; seq_id <= stop; seq_id++)
    {
      seq_ids.push_back(seq_id);
    }
  }
  else
  {
    while(optind < argc)
    {
      seq_ids.push_back(std::stoul(argv[optind]));
      optind++;
    }
  }

  Version::print(std::cout, tool_name);

  printHeader("Input"); std::cout << base_name << std::endl;
  printHeader("Output"); std::cout << output << std::endl;
  printHeader("Sequences"); std::cout << seq_ids.size() << std::endl;
  printHeader("Chunk size"); std::cout << chunk_size << std::endl;
  std::cout << std::endl;

  double start = readTimer();

  DynamicGBWT index;
  if(!sdsl::load_from_file(index, base_name + DynamicGBWT::EXTENSION))
  {
    std::cerr << "remove_seq: Cannot load the index from " << (base_name + DynamicGBWT::EXTENSION) << std::endl;
    std::exit(EXIT_FAILURE);
  }
  printStatistics(index, base_name);

  size_type total_length = index.remove(seq_ids, chunk_size);
  if(total_length > 0)
  {
    if(!sdsl::store_to_file(index, output + DynamicGBWT::EXTENSION))
    {
      std::cerr << "remove_seq: Cannot write the index to " << (output + DynamicGBWT::EXTENSION) << std::endl;
      std::exit(EXIT_FAILURE);
    }
    printStatistics(index, output);
  }

  double seconds = readTimer() - start;

  std::cout << "Removed " << total_length << " nodes in " << seconds << " seconds ("
            << (total_length / seconds) << " nodes/second)" << std::endl;
  std::cout << "Memory usage " << inGigabytes(memoryUsage()) << " GB" << std::endl;
  std::cout << std::endl;

  return 0;
}

//------------------------------------------------------------------------------

void
printUsage(int exit_code)
{
  Version::print(std::cerr, tool_name);

  std::cerr << "Usage: remove_seq [options] base_name seq1 [seq2 ...]" << std::endl;
  std::cerr << std::endl;
  std::cerr << "  -c N  Build the RA in chunks of N sequences per thread (default: " << DynamicGBWT::REMOVE_CHUNK_SIZE << ")" << std::endl;
  std::cerr << "  -o X  Use X as the base name for output" << std::endl;
  std::cerr << "  -r    Remove a range of sequences (inclusive; requires 2 sequence ids)" << std::endl;
  std::cerr << std::endl;

  std::exit(exit_code);
}

//------------------------------------------------------------------------------
<commit_msg>Remove sequences by sample/contig in remove_seq<commit_after>/*
  Copyright (c) 2018, 2019 Jouni Siren

  Author: Jouni Siren <jouni.siren@iki.fi>

  Permission is hereby granted, free of charge, to any person obtaining a copy
  of this software and associated documentation files (the "Software"), to deal
  in the Software without restriction, including without limitation the rights
  to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
  copies of the Software, and to permit persons to whom the Software is
  furnished to do so, subject to the following conditions:

  The above copyright notice and this permission notice shall be included in all
  copies or substantial portions of the Software.

  THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
  IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
  FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
  AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
  LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
  OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
  SOFTWARE.
*/

#include <string>
#include <unistd.h>

#include <gbwt/dynamic_gbwt.h>

using namespace gbwt;

//------------------------------------------------------------------------------

const std::string tool_name = "GBWT sequence removal";

void printUsage(int exit_code = EXIT_SUCCESS);

//------------------------------------------------------------------------------

int
main(int argc, char** argv)
{
  if(argc < 3) { printUsage(); }

  // Parse command line options.
  int c = 0;
  std::string output;
  size_type chunk_size = DynamicGBWT::REMOVE_CHUNK_SIZE;
  bool range = false, sample = false, contig = false;
  while((c = getopt(argc, argv, "c:o:rSC")) != -1)
  {
    switch(c)
    {
    case 'c':
      chunk_size = std::max(1ul, std::stoul(optarg));
      break;
    case 'o':
      output = optarg; break;
    case 'r':
      range = true; break;
    case 'S':
      sample = true; break;
    case 'C':
      contig = true; break;
    case '?':
      std::exit(EXIT_FAILURE);
    default:
      std::exit(EXIT_FAILURE);
    }
  }

  // Check command line options.
  if(optind + 1 >= argc) { printUsage(EXIT_FAILURE); }
  if((range & sample) || (range & contig) || (sample & contig))
  {
    std::cerr << "remove_seq: Options -r, -S, and -C are mutually exclusive" << std::endl;
    std::exit(EXIT_FAILURE);
  }
  std::string base_name = argv[optind]; optind++;
  std::string key;
  if(output.empty()) { output = base_name; }
  std::vector<size_type> seq_ids;
  if(range)
  {
    if(argc != optind + 2) { printUsage(EXIT_FAILURE); }
    size_type start = std::stoul(argv[optind]); optind++;
    size_type stop = std::stoul(argv[optind]); optind++;
    if(stop < start) { printUsage(EXIT_FAILURE); }
    for(size_type seq_id = start; seq_id <= stop; seq_id++)
    {
      seq_ids.push_back(seq_id);
    }
  }
  else if(sample || contig)
  {
    if(argc != optind + 1) { printUsage(EXIT_FAILURE); }
    key = argv[optind];
  }
  else
  {
    while(optind < argc)
    {
      seq_ids.push_back(std::stoul(argv[optind]));
      optind++;
    }
  }

  // Initial output.
  Version::print(std::cout, tool_name);
  printHeader("Input"); std::cout << base_name << std::endl;
  printHeader("Output"); std::cout << output << std::endl;
  if(range)
  {
    printHeader("Range"); std::cout << range_type(seq_ids.front(), seq_ids.back()) << std::endl;
  }
  else if(sample)
  {
    printHeader("Sample"); std::cout << key << std::endl;
  }
  else if(contig)
  {
    printHeader("Contig"); std::cout << key << std::endl;
  }
  else
  {
    printHeader("Sequences"); std::cout << seq_ids.size() << std::endl;
  }
  printHeader("Chunk size"); std::cout << chunk_size << std::endl;
  std::cout << std::endl;

  double start = readTimer();

  // Load index.
  DynamicGBWT index;
  if(!sdsl::load_from_file(index, base_name + DynamicGBWT::EXTENSION))
  {
    std::cerr << "remove_seq: Cannot load the index from " << (base_name + DynamicGBWT::EXTENSION) << std::endl;
    std::exit(EXIT_FAILURE);
  }
  printStatistics(index, base_name);

  // Handle metadata.
  if(sample)
  {
    if(!(index.hasMetadata()) || !(index.metadata.hasSampleNames()) || !(index.metadata.hasPathNames()))
    {
      std::cerr << "remove_seq: Option -S requires sample and path names" << std::endl;
      std::exit(EXIT_FAILURE);
    }
    size_type sample_id = index.metadata.sample(key);
    seq_ids = index.metadata.removeSample(sample_id);
    if(seq_ids.empty())
    {
      std::cerr << "remove_seq: No sequences for sample " << key << std::endl;
      std::exit(EXIT_FAILURE);
    }
  }
  else if(contig)
  {
    if(!(index.hasMetadata()) || !(index.metadata.hasContigNames()) || !(index.metadata.hasPathNames()))
    {
      std::cerr << "remove_seq: Option -C requires contig and path names" << std::endl;
      std::exit(EXIT_FAILURE);
    }
    size_type contig_id = index.metadata.contig(key);
    seq_ids = index.metadata.removeContig(contig_id);
    if(seq_ids.empty())
    {
      std::cerr << "remove_seq: No sequences for contig " << key << std::endl;
      std::exit(EXIT_FAILURE);
    }
  }
  else
  {
    if(index.hasMetadata() && index.metadata.hasPathNames())
    {
      std::cerr << "remove_seq: Removing arbitrary sequences would leave the metadata inconsistent" << std::endl;
      std::exit(EXIT_FAILURE);
    }
  }

  // Remove the sequences.
  size_type total_length = index.remove(seq_ids, chunk_size);
  if(total_length > 0)
  {
    if(!sdsl::store_to_file(index, output + DynamicGBWT::EXTENSION))
    {
      std::cerr << "remove_seq: Cannot write the index to " << (output + DynamicGBWT::EXTENSION) << std::endl;
      std::exit(EXIT_FAILURE);
    }
    printStatistics(index, output);
  }

  double seconds = readTimer() - start;

  std::cout << "Removed " << total_length << " nodes in " << seconds << " seconds ("
            << (total_length / seconds) << " nodes/second)" << std::endl;
  std::cout << "Memory usage " << inGigabytes(memoryUsage()) << " GB" << std::endl;
  std::cout << std::endl;

  return 0;
}

//------------------------------------------------------------------------------

void
printUsage(int exit_code)
{
  Version::print(std::cerr, tool_name);

  std::cerr << "Usage: remove_seq [options] base_name seq1 [seq2 ...]" << std::endl;
  std::cerr << std::endl;
  std::cerr << "  -c N  Build the RA in chunks of N sequences per thread (default: " << DynamicGBWT::REMOVE_CHUNK_SIZE << ")" << std::endl;
  std::cerr << "  -o X  Use X as the base name for output" << std::endl;
  std::cerr << "  -r    Remove a range of sequences (inclusive; requires 2 sequence ids)" << std::endl;
  std::cerr << "  -S    Remove all sequences for the sample with name seq1 (cannot have seq2)" << std::endl;
  std::cerr << "  -C    Remove all sequences for the contig with name seq1 (cannot have seq2)" << std::endl;
  std::cerr << std::endl;

  std::exit(exit_code);
}

//------------------------------------------------------------------------------
<|endoftext|>


<commit_before>#include "ylikuutio_string.hpp"

// Include standard headers
#include <cstdio>   // std::FILE, std::fclose, std::fopen, std::fread, std::getchar, std::printf etc.
#include <cstring>  // std::memcmp, std::strcmp, std::strlen, std::strncmp
#include <iostream> // std::cout, std::cin, std::cerr
#include <list>     // std::list
#include <sstream>  // std::stringstream
#include <string>   // std::string
#include <vector>   // std::vector

namespace string
{
    bool check_and_report_if_some_string_matches(const char* file_base_pointer, char* file_data_pointer, std::vector<std::string> identifier_strings_vector)
    {
        for (std::string identifier_string : identifier_strings_vector)
        {
            const char* identifier_string_char = identifier_string.c_str();

            if (std::strncmp(file_data_pointer, identifier_string_char, std::strlen(identifier_string_char)) == 0)
            {
                const char* identifier_string_char = identifier_string.c_str();
                uint64_t offset = (uint64_t) file_data_pointer - (uint64_t) file_base_pointer;
                std::printf("%s found at file offset 0x%llu (memory address 0x%llu).\n", identifier_string_char, offset, (uint64_t) file_data_pointer);
                return true;
            }
        }
        return false;
    }

    void extract_string(char* dest_mem_pointer, char* &src_mem_pointer, char* char_end_string)
    {
        while (std::strncmp(src_mem_pointer, char_end_string, std::strlen(char_end_string)) != 0)
        {
            strncpy(dest_mem_pointer++, src_mem_pointer++, 1);
        }
        *dest_mem_pointer = '\0';
    }

    void extract_string_with_several_endings(char* dest_mem_pointer, char*& src_mem_pointer, char* char_end_string)
    {
        // This function copies characters from `src_mem_pointer` until a character matches.

        while (true)
        {
            uint32_t n_of_ending_characters = std::strlen(char_end_string);
            char* end_char_pointer;
            end_char_pointer = char_end_string;

            // Check if current character is any of the ending characters.
            while (*end_char_pointer != '\0')
            {
                if (std::strncmp(src_mem_pointer, end_char_pointer, 1) == 0)
                {
                    *dest_mem_pointer = '\0';
                    return;
                }
                end_char_pointer++;
            }

            // OK, current character is not any of the ending characters.
            // Copy it and advance the pointers accordingly.
            strncpy(dest_mem_pointer++, src_mem_pointer++, 1);
        }
    }

    int32_t extract_int32_t_value_from_string(char*& data_pointer, char* char_end_string, const char* description)
    {
        char char_number_buffer[1024]; // FIXME: risk of buffer overflow.
        char* dest_mem_pointer;
        dest_mem_pointer = char_number_buffer;
        string::extract_string_with_several_endings(dest_mem_pointer, ++data_pointer, char_end_string);
        uint32_t value = std::atoi(dest_mem_pointer);
        std::printf("%s: %d\n", description, value);
        return value;
    }

    float extract_float_value_from_string(char*& data_pointer, char* char_end_string, const char* description)
    {
        char char_number_buffer[1024]; // FIXME: risk of buffer overflow.
        char* dest_mem_pointer;
        dest_mem_pointer = char_number_buffer;
        string::extract_string_with_several_endings(dest_mem_pointer, ++data_pointer, char_end_string);
        float value = std::atof(dest_mem_pointer);
        std::printf("%s: %f\n", description, value);
        return value;
    }

    int32_t extract_unicode_value_from_string(const char*& unicode_char_pointer)
    {
        if (*unicode_char_pointer == '\0')
        {
            unicode_char_pointer++;
            std::cerr << "Error: Unicode can not begin with \\0!\n";
            return 0xdfff; // invalid unicode!
        }

        if (*unicode_char_pointer != '&')
        {
            // it's just a character, so return its value,
            // and advance to the next character.
            return (int32_t) *unicode_char_pointer++;
        }

        if (*++unicode_char_pointer != '#')
        {
            // not valid format, must begin `"&#x"`.
            unicode_char_pointer++;
            std::cerr << "Error: Unicode string format not supported!\n";
            return 0xdfff; // invalid unicode!
        }

        if (*++unicode_char_pointer != 'x')
        {
            // not valid format, must begin `"&#x"`.
            unicode_char_pointer++;
            std::cerr << "Error: Unicode string format not supported!\n";
            return 0xdfff; // invalid unicode!
        }

        // valid format.
        std::string hex_string;

        // unicode string beginning with '&'
        while (*++unicode_char_pointer != ';')
        {
            if (*unicode_char_pointer == '\0')
            {
                std::cerr << "Error: Null character \\0 reached before end of Unicode string!\n";
                return 0xdfff; // invalid unicode!
            }

            char current_char = *unicode_char_pointer;
            hex_string.append(unicode_char_pointer);
        }

        // Advance to the next character.
        unicode_char_pointer++;

        // convert hexadecimal string to signed integer.
        // http://stackoverflow.com/questions/1070497/c-convert-hex-string-to-signed-integer/1070499#1070499
        uint32_t unicode_value;
        std::stringstream unicode_stringstream;
        unicode_stringstream << std::hex << hex_string;
        unicode_stringstream >> unicode_value;
        return unicode_value;
    }

    std::string convert_std_list_char_to_std_string(const std::list<char>& std_list_char)
    {
        std::string my_string;

        for (std::list<char>::const_iterator it = std_list_char.begin(); it != std_list_char.end(); it++)
        {
            my_string.push_back(*it);
        }

        return my_string;
    }

    std::string convert_std_list_char_to_std_string(const std::list<char>& std_list_char, uint32_t first_line_length, uint32_t line_length)
    {
        std::string my_string;
        uint32_t remaining_characters_on_this_line = first_line_length;

        for (std::list<char>::const_iterator it = std_list_char.begin(); it != std_list_char.end(); it++)
        {
            if (remaining_characters_on_this_line == 0)
            {
                my_string.push_back('\\');
                my_string.push_back('n');
                remaining_characters_on_this_line = line_length;
            }
            my_string.push_back(*it);
            remaining_characters_on_this_line--;
        }
        return my_string;
    }
}
<commit_msg>If `description` is `nullptr`, do not `printf` anything.<commit_after>#include "ylikuutio_string.hpp"

// Include standard headers
#include <cstdio>   // std::FILE, std::fclose, std::fopen, std::fread, std::getchar, std::printf etc.
#include <cstring>  // std::memcmp, std::strcmp, std::strlen, std::strncmp
#include <iostream> // std::cout, std::cin, std::cerr
#include <list>     // std::list
#include <sstream>  // std::stringstream
#include <string>   // std::string
#include <vector>   // std::vector

namespace string
{
    bool check_and_report_if_some_string_matches(const char* file_base_pointer, char* file_data_pointer, std::vector<std::string> identifier_strings_vector)
    {
        for (std::string identifier_string : identifier_strings_vector)
        {
            const char* identifier_string_char = identifier_string.c_str();

            if (std::strncmp(file_data_pointer, identifier_string_char, std::strlen(identifier_string_char)) == 0)
            {
                const char* identifier_string_char = identifier_string.c_str();
                uint64_t offset = (uint64_t) file_data_pointer - (uint64_t) file_base_pointer;
                std::printf("%s found at file offset 0x%llu (memory address 0x%llu).\n", identifier_string_char, offset, (uint64_t) file_data_pointer);
                return true;
            }
        }
        return false;
    }

    void extract_string(char* dest_mem_pointer, char* &src_mem_pointer, char* char_end_string)
    {
        while (std::strncmp(src_mem_pointer, char_end_string, std::strlen(char_end_string)) != 0)
        {
            strncpy(dest_mem_pointer++, src_mem_pointer++, 1);
        }
        *dest_mem_pointer = '\0';
    }

    void extract_string_with_several_endings(char* dest_mem_pointer, char*& src_mem_pointer, char* char_end_string)
    {
        // This function copies characters from `src_mem_pointer` until a character matches.

        while (true)
        {
            uint32_t n_of_ending_characters = std::strlen(char_end_string);
            char* end_char_pointer;
            end_char_pointer = char_end_string;

            // Check if current character is any of the ending characters.
            while (*end_char_pointer != '\0')
            {
                if (std::strncmp(src_mem_pointer, end_char_pointer, 1) == 0)
                {
                    *dest_mem_pointer = '\0';
                    return;
                }
                end_char_pointer++;
            }

            // OK, current character is not any of the ending characters.
            // Copy it and advance the pointers accordingly.
            strncpy(dest_mem_pointer++, src_mem_pointer++, 1);
        }
    }

    int32_t extract_int32_t_value_from_string(char*& data_pointer, char* char_end_string, const char* description)
    {
        char char_number_buffer[1024]; // FIXME: risk of buffer overflow.
        char* dest_mem_pointer;
        dest_mem_pointer = char_number_buffer;
        string::extract_string_with_several_endings(dest_mem_pointer, ++data_pointer, char_end_string);
        uint32_t value = std::atoi(dest_mem_pointer);
        if (description != nullptr)
        {
            std::printf("%s: %d\n", description, value);
        }
        return value;
    }

    float extract_float_value_from_string(char*& data_pointer, char* char_end_string, const char* description)
    {
        char char_number_buffer[1024]; // FIXME: risk of buffer overflow.
        char* dest_mem_pointer;
        dest_mem_pointer = char_number_buffer;
        string::extract_string_with_several_endings(dest_mem_pointer, ++data_pointer, char_end_string);
        float value = std::atof(dest_mem_pointer);

        if (description != nullptr)
        {
            std::printf("%s: %f\n", description, value);
        }
        return value;
    }

    int32_t extract_unicode_value_from_string(const char*& unicode_char_pointer)
    {
        if (*unicode_char_pointer == '\0')
        {
            unicode_char_pointer++;
            std::cerr << "Error: Unicode can not begin with \\0!\n";
            return 0xdfff; // invalid unicode!
        }

        if (*unicode_char_pointer != '&')
        {
            // it's just a character, so return its value,
            // and advance to the next character.
            return (int32_t) *unicode_char_pointer++;
        }

        if (*++unicode_char_pointer != '#')
        {
            // not valid format, must begin `"&#x"`.
            unicode_char_pointer++;
            std::cerr << "Error: Unicode string format not supported!\n";
            return 0xdfff; // invalid unicode!
        }

        if (*++unicode_char_pointer != 'x')
        {
            // not valid format, must begin `"&#x"`.
            unicode_char_pointer++;
            std::cerr << "Error: Unicode string format not supported!\n";
            return 0xdfff; // invalid unicode!
        }

        // valid format.
        std::string hex_string;

        // unicode string beginning with '&'
        while (*++unicode_char_pointer != ';')
        {
            if (*unicode_char_pointer == '\0')
            {
                std::cerr << "Error: Null character \\0 reached before end of Unicode string!\n";
                return 0xdfff; // invalid unicode!
            }

            char current_char = *unicode_char_pointer;
            hex_string.append(unicode_char_pointer);
        }

        // Advance to the next character.
        unicode_char_pointer++;

        // convert hexadecimal string to signed integer.
        // http://stackoverflow.com/questions/1070497/c-convert-hex-string-to-signed-integer/1070499#1070499
        uint32_t unicode_value;
        std::stringstream unicode_stringstream;
        unicode_stringstream << std::hex << hex_string;
        unicode_stringstream >> unicode_value;
        return unicode_value;
    }

    std::string convert_std_list_char_to_std_string(const std::list<char>& std_list_char)
    {
        std::string my_string;

        for (std::list<char>::const_iterator it = std_list_char.begin(); it != std_list_char.end(); it++)
        {
            my_string.push_back(*it);
        }

        return my_string;
    }

    std::string convert_std_list_char_to_std_string(const std::list<char>& std_list_char, uint32_t first_line_length, uint32_t line_length)
    {
        std::string my_string;
        uint32_t remaining_characters_on_this_line = first_line_length;

        for (std::list<char>::const_iterator it = std_list_char.begin(); it != std_list_char.end(); it++)
        {
            if (remaining_characters_on_this_line == 0)
            {
                my_string.push_back('\\');
                my_string.push_back('n');
                remaining_characters_on_this_line = line_length;
            }
            my_string.push_back(*it);
            remaining_characters_on_this_line--;
        }
        return my_string;
    }
}
<|endoftext|>


<commit_before>// Copyright (c) 2009 The Chromium Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#include "app/gtk_dnd_util.h"

#include "base/logging.h"
#include "base/pickle.h"
#include "base/utf_string_conversions.h"
#include "googleurl/src/gurl.h"

static const int kBitsPerByte = 8;

using WebKit::WebDragOperationsMask;
using WebKit::WebDragOperation;
using WebKit::WebDragOperationNone;
using WebKit::WebDragOperationCopy;
using WebKit::WebDragOperationLink;
using WebKit::WebDragOperationMove;

namespace {

void AddTargetToList(GtkTargetList* targets, int target_code) {
  switch (target_code) {
    case gtk_dnd_util::TEXT_PLAIN:
      gtk_target_list_add_text_targets(targets, gtk_dnd_util::TEXT_PLAIN);
      break;

    case gtk_dnd_util::TEXT_URI_LIST:
      gtk_target_list_add_uri_targets(targets, gtk_dnd_util::TEXT_URI_LIST);
      break;

    case gtk_dnd_util::TEXT_HTML:
      gtk_target_list_add(
          targets, gtk_dnd_util::GetAtomForTarget(gtk_dnd_util::TEXT_PLAIN),
          0, gtk_dnd_util::TEXT_HTML);
      break;

    case gtk_dnd_util::NETSCAPE_URL:
      gtk_target_list_add(targets,
          gtk_dnd_util::GetAtomForTarget(gtk_dnd_util::NETSCAPE_URL),
          0, gtk_dnd_util::NETSCAPE_URL);
      break;

    case gtk_dnd_util::CHROME_TAB:
    case gtk_dnd_util::CHROME_BOOKMARK_ITEM:
    case gtk_dnd_util::CHROME_NAMED_URL:
      gtk_target_list_add(targets, gtk_dnd_util::GetAtomForTarget(target_code),
                          GTK_TARGET_SAME_APP, target_code);
      break;

    case gtk_dnd_util::DIRECT_SAVE_FILE:
      gtk_target_list_add(targets,
          gtk_dnd_util::GetAtomForTarget(gtk_dnd_util::DIRECT_SAVE_FILE),
          0, gtk_dnd_util::DIRECT_SAVE_FILE);
      break;

    default:
      NOTREACHED() << " Unexpected target code: " << target_code;
  }
}

}  // namespace

namespace gtk_dnd_util {

GdkAtom GetAtomForTarget(int target) {
  switch (target) {
    case CHROME_TAB:
      static GdkAtom tab_atom = gdk_atom_intern(
          const_cast<char*>("application/x-chrome-tab"), false);
      return tab_atom;

    case TEXT_HTML:
      static GdkAtom html_atom = gdk_atom_intern(
          const_cast<char*>("text/html"), false);
      return html_atom;

    case CHROME_BOOKMARK_ITEM:
      static GdkAtom bookmark_atom = gdk_atom_intern(
          const_cast<char*>("application/x-chrome-bookmark-item"), false);
      return bookmark_atom;

    case TEXT_PLAIN:
      static GdkAtom text_atom = gdk_atom_intern(
          const_cast<char*>("text/plain;charset=utf-8"), false);
      return text_atom;

    case TEXT_URI_LIST:
      static GdkAtom uris_atom = gdk_atom_intern(
          const_cast<char*>("text/uri-list"), false);
      return uris_atom;

    case CHROME_NAMED_URL:
      static GdkAtom named_url = gdk_atom_intern(
          const_cast<char*>("application/x-chrome-named-url"), false);
      return named_url;

    case NETSCAPE_URL:
      static GdkAtom netscape_url = gdk_atom_intern(
          const_cast<char*>("_NETSCAPE_URL"), false);
      return netscape_url;

    case TEXT_PLAIN_NO_CHARSET:
      static GdkAtom text_no_charset_atom = gdk_atom_intern(
          const_cast<char*>("text/plain"), false);
      return text_no_charset_atom;

    case DIRECT_SAVE_FILE:
      static GdkAtom xds_atom = gdk_atom_intern(
          const_cast<char*>("XdndDirectSave0"), false);
      return xds_atom;

    default:
      NOTREACHED();
  }

  return NULL;
}

GtkTargetList* GetTargetListFromCodeMask(int code_mask) {
  GtkTargetList* targets = gtk_target_list_new(NULL, 0);

  for (size_t i = 1; i < INVALID_TARGET; i = i << 1) {
    if (i == CHROME_WEBDROP_FILE_CONTENTS)
      continue;

    if (i & code_mask)
      AddTargetToList(targets, i);
  }

  return targets;
}

void SetSourceTargetListFromCodeMask(GtkWidget* source, int code_mask) {
  GtkTargetList* targets = GetTargetListFromCodeMask(code_mask);
  gtk_drag_source_set_target_list(source, targets);
  gtk_target_list_unref(targets);
}

void SetDestTargetList(GtkWidget* dest, const int* target_codes) {
  GtkTargetList* targets = gtk_target_list_new(NULL, 0);

  for (size_t i = 0; target_codes[i] != -1; ++i) {
    AddTargetToList(targets, target_codes[i]);
  }

  gtk_drag_dest_set_target_list(dest, targets);
  gtk_target_list_unref(targets);
}

void WriteURLWithName(GtkSelectionData* selection_data,
                      const GURL& url,
                      const string16& title,
                      int type) {
  switch (type) {
    case TEXT_PLAIN: {
      gtk_selection_data_set_text(selection_data, url.spec().c_str(),
                                  url.spec().length());
      break;
    }
    case TEXT_URI_LIST: {
      gchar* uri_array[2];
      uri_array[0] = strdup(url.spec().c_str());
      uri_array[1] = NULL;
      gtk_selection_data_set_uris(selection_data, uri_array);
      free(uri_array[0]);
      break;
    }
    case CHROME_NAMED_URL: {
      Pickle pickle;
      pickle.WriteString(UTF16ToUTF8(title));
      pickle.WriteString(url.spec());
      gtk_selection_data_set(
          selection_data,
          GetAtomForTarget(gtk_dnd_util::CHROME_NAMED_URL),
          kBitsPerByte,
          reinterpret_cast<const guchar*>(pickle.data()),
          pickle.size());
      break;
    }
    case NETSCAPE_URL: {
      // _NETSCAPE_URL format is URL + \n + title.
      std::string utf8_text = url.spec() + "\n" + UTF16ToUTF8(title);
      gtk_selection_data_set(selection_data,
                             selection_data->target,
                             kBitsPerByte,
                             reinterpret_cast<const guchar*>(utf8_text.c_str()),
                             utf8_text.length());
      break;
    }

    default: {
      NOTREACHED();
      break;
    }
  }
}

bool ExtractNamedURL(GtkSelectionData* selection_data,
                     GURL* url,
                     string16* title) {
  Pickle data(reinterpret_cast<char*>(selection_data->data),
              selection_data->length);
  void* iter = NULL;
  std::string title_utf8, url_utf8;
  if (!data.ReadString(&iter, &title_utf8) ||
      !data.ReadString(&iter, &url_utf8)) {
    return false;
  }

  GURL gurl(url_utf8);
  if (!gurl.is_valid())
    return false;

  *url = gurl;
  *title = UTF8ToUTF16(title_utf8);
  return true;
}

bool ExtractURIList(GtkSelectionData* selection_data, std::vector<GURL>* urls) {
  gchar** uris = gtk_selection_data_get_uris(selection_data);
  if (!uris)
    return false;

  for (size_t i = 0; uris[i] != NULL; ++i) {
    GURL url(uris[i]);
    if (url.is_valid())
      urls->push_back(url);
  }

  g_strfreev(uris);
  return true;
}

GdkDragAction WebDragOpToGdkDragAction(WebDragOperationsMask op) {
  GdkDragAction action = static_cast<GdkDragAction>(0);
  if (op & WebDragOperationCopy)
    action = static_cast<GdkDragAction>(action | GDK_ACTION_COPY);
  if (op & WebDragOperationLink)
    action = static_cast<GdkDragAction>(action | GDK_ACTION_LINK);
  if (op & WebDragOperationMove)
    action = static_cast<GdkDragAction>(action | GDK_ACTION_MOVE);
  return action;
}

WebDragOperationsMask GdkDragActionToWebDragOp(GdkDragAction action) {
  WebDragOperationsMask op = WebDragOperationNone;
  if (action & GDK_ACTION_COPY)
    op = static_cast<WebDragOperationsMask>(op | WebDragOperationCopy);
  if (action & GDK_ACTION_LINK)
    op = static_cast<WebDragOperationsMask>(op | WebDragOperationLink);
  if (action & GDK_ACTION_MOVE)
    op = static_cast<WebDragOperationsMask>(op | WebDragOperationMove);
  return op;
}

}  // namespace gtk_dnd_util
<commit_msg>We DnD shouldn't provide text/html if it claims to provide text/text/plain;charset=utf-8.<commit_after>// Copyright (c) 2009 The Chromium Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#include "app/gtk_dnd_util.h"

#include "base/logging.h"
#include "base/pickle.h"
#include "base/utf_string_conversions.h"
#include "googleurl/src/gurl.h"

static const int kBitsPerByte = 8;

using WebKit::WebDragOperationsMask;
using WebKit::WebDragOperation;
using WebKit::WebDragOperationNone;
using WebKit::WebDragOperationCopy;
using WebKit::WebDragOperationLink;
using WebKit::WebDragOperationMove;

namespace {

void AddTargetToList(GtkTargetList* targets, int target_code) {
  switch (target_code) {
    case gtk_dnd_util::TEXT_PLAIN:
      gtk_target_list_add_text_targets(targets, gtk_dnd_util::TEXT_PLAIN);
      break;

    case gtk_dnd_util::TEXT_URI_LIST:
      gtk_target_list_add_uri_targets(targets, gtk_dnd_util::TEXT_URI_LIST);
      break;

    case gtk_dnd_util::TEXT_HTML:
      gtk_target_list_add(
          targets, gtk_dnd_util::GetAtomForTarget(gtk_dnd_util::TEXT_HTML),
          0, gtk_dnd_util::TEXT_HTML);
      break;

    case gtk_dnd_util::NETSCAPE_URL:
      gtk_target_list_add(targets,
          gtk_dnd_util::GetAtomForTarget(gtk_dnd_util::NETSCAPE_URL),
          0, gtk_dnd_util::NETSCAPE_URL);
      break;

    case gtk_dnd_util::CHROME_TAB:
    case gtk_dnd_util::CHROME_BOOKMARK_ITEM:
    case gtk_dnd_util::CHROME_NAMED_URL:
      gtk_target_list_add(targets, gtk_dnd_util::GetAtomForTarget(target_code),
                          GTK_TARGET_SAME_APP, target_code);
      break;

    case gtk_dnd_util::DIRECT_SAVE_FILE:
      gtk_target_list_add(targets,
          gtk_dnd_util::GetAtomForTarget(gtk_dnd_util::DIRECT_SAVE_FILE),
          0, gtk_dnd_util::DIRECT_SAVE_FILE);
      break;

    default:
      NOTREACHED() << " Unexpected target code: " << target_code;
  }
}

}  // namespace

namespace gtk_dnd_util {

GdkAtom GetAtomForTarget(int target) {
  switch (target) {
    case CHROME_TAB:
      static GdkAtom tab_atom = gdk_atom_intern(
          const_cast<char*>("application/x-chrome-tab"), false);
      return tab_atom;

    case TEXT_HTML:
      static GdkAtom html_atom = gdk_atom_intern(
          const_cast<char*>("text/html"), false);
      return html_atom;

    case CHROME_BOOKMARK_ITEM:
      static GdkAtom bookmark_atom = gdk_atom_intern(
          const_cast<char*>("application/x-chrome-bookmark-item"), false);
      return bookmark_atom;

    case TEXT_PLAIN:
      static GdkAtom text_atom = gdk_atom_intern(
          const_cast<char*>("text/plain;charset=utf-8"), false);
      return text_atom;

    case TEXT_URI_LIST:
      static GdkAtom uris_atom = gdk_atom_intern(
          const_cast<char*>("text/uri-list"), false);
      return uris_atom;

    case CHROME_NAMED_URL:
      static GdkAtom named_url = gdk_atom_intern(
          const_cast<char*>("application/x-chrome-named-url"), false);
      return named_url;

    case NETSCAPE_URL:
      static GdkAtom netscape_url = gdk_atom_intern(
          const_cast<char*>("_NETSCAPE_URL"), false);
      return netscape_url;

    case TEXT_PLAIN_NO_CHARSET:
      static GdkAtom text_no_charset_atom = gdk_atom_intern(
          const_cast<char*>("text/plain"), false);
      return text_no_charset_atom;

    case DIRECT_SAVE_FILE:
      static GdkAtom xds_atom = gdk_atom_intern(
          const_cast<char*>("XdndDirectSave0"), false);
      return xds_atom;

    default:
      NOTREACHED();
  }

  return NULL;
}

GtkTargetList* GetTargetListFromCodeMask(int code_mask) {
  GtkTargetList* targets = gtk_target_list_new(NULL, 0);

  for (size_t i = 1; i < INVALID_TARGET; i = i << 1) {
    if (i == CHROME_WEBDROP_FILE_CONTENTS)
      continue;

    if (i & code_mask)
      AddTargetToList(targets, i);
  }

  return targets;
}

void SetSourceTargetListFromCodeMask(GtkWidget* source, int code_mask) {
  GtkTargetList* targets = GetTargetListFromCodeMask(code_mask);
  gtk_drag_source_set_target_list(source, targets);
  gtk_target_list_unref(targets);
}

void SetDestTargetList(GtkWidget* dest, const int* target_codes) {
  GtkTargetList* targets = gtk_target_list_new(NULL, 0);

  for (size_t i = 0; target_codes[i] != -1; ++i) {
    AddTargetToList(targets, target_codes[i]);
  }

  gtk_drag_dest_set_target_list(dest, targets);
  gtk_target_list_unref(targets);
}

void WriteURLWithName(GtkSelectionData* selection_data,
                      const GURL& url,
                      const string16& title,
                      int type) {
  switch (type) {
    case TEXT_PLAIN: {
      gtk_selection_data_set_text(selection_data, url.spec().c_str(),
                                  url.spec().length());
      break;
    }
    case TEXT_URI_LIST: {
      gchar* uri_array[2];
      uri_array[0] = strdup(url.spec().c_str());
      uri_array[1] = NULL;
      gtk_selection_data_set_uris(selection_data, uri_array);
      free(uri_array[0]);
      break;
    }
    case CHROME_NAMED_URL: {
      Pickle pickle;
      pickle.WriteString(UTF16ToUTF8(title));
      pickle.WriteString(url.spec());
      gtk_selection_data_set(
          selection_data,
          GetAtomForTarget(gtk_dnd_util::CHROME_NAMED_URL),
          kBitsPerByte,
          reinterpret_cast<const guchar*>(pickle.data()),
          pickle.size());
      break;
    }
    case NETSCAPE_URL: {
      // _NETSCAPE_URL format is URL + \n + title.
      std::string utf8_text = url.spec() + "\n" + UTF16ToUTF8(title);
      gtk_selection_data_set(selection_data,
                             selection_data->target,
                             kBitsPerByte,
                             reinterpret_cast<const guchar*>(utf8_text.c_str()),
                             utf8_text.length());
      break;
    }

    default: {
      NOTREACHED();
      break;
    }
  }
}

bool ExtractNamedURL(GtkSelectionData* selection_data,
                     GURL* url,
                     string16* title) {
  Pickle data(reinterpret_cast<char*>(selection_data->data),
              selection_data->length);
  void* iter = NULL;
  std::string title_utf8, url_utf8;
  if (!data.ReadString(&iter, &title_utf8) ||
      !data.ReadString(&iter, &url_utf8)) {
    return false;
  }

  GURL gurl(url_utf8);
  if (!gurl.is_valid())
    return false;

  *url = gurl;
  *title = UTF8ToUTF16(title_utf8);
  return true;
}

bool ExtractURIList(GtkSelectionData* selection_data, std::vector<GURL>* urls) {
  gchar** uris = gtk_selection_data_get_uris(selection_data);
  if (!uris)
    return false;

  for (size_t i = 0; uris[i] != NULL; ++i) {
    GURL url(uris[i]);
    if (url.is_valid())
      urls->push_back(url);
  }

  g_strfreev(uris);
  return true;
}

GdkDragAction WebDragOpToGdkDragAction(WebDragOperationsMask op) {
  GdkDragAction action = static_cast<GdkDragAction>(0);
  if (op & WebDragOperationCopy)
    action = static_cast<GdkDragAction>(action | GDK_ACTION_COPY);
  if (op & WebDragOperationLink)
    action = static_cast<GdkDragAction>(action | GDK_ACTION_LINK);
  if (op & WebDragOperationMove)
    action = static_cast<GdkDragAction>(action | GDK_ACTION_MOVE);
  return action;
}

WebDragOperationsMask GdkDragActionToWebDragOp(GdkDragAction action) {
  WebDragOperationsMask op = WebDragOperationNone;
  if (action & GDK_ACTION_COPY)
    op = static_cast<WebDragOperationsMask>(op | WebDragOperationCopy);
  if (action & GDK_ACTION_LINK)
    op = static_cast<WebDragOperationsMask>(op | WebDragOperationLink);
  if (action & GDK_ACTION_MOVE)
    op = static_cast<WebDragOperationsMask>(op | WebDragOperationMove);
  return op;
}

}  // namespace gtk_dnd_util
<|endoftext|>


<commit_before>#include "types.h"
#include "amd64.h"
#include "mmu.h"
#include "cpu.hh"
#include "kernel.hh"
#include "bits.hh"
#include "spinlock.h"
#include "kalloc.hh"
#include "queue.h"
#include "condvar.h"
#include "proc.hh"
#include "vm.hh"
#include <stddef.h>

using namespace std;

static pgmap*
descend(pgmap *dir, u64 va, u64 flags, int create, int level)
{
  atomic<pme_t> *entryp;
  pme_t entry;
  pgmap *next;

retry:
  entryp = &dir->e[PX(level, va)];
  entry = entryp->load();
  if (entry & PTE_P) {
    next = (pgmap*) p2v(PTE_ADDR(entry));
  } else {
    if (!create)
      return NULL;
    next = (pgmap*) kalloc();
    if (!next)
      return NULL;
    memset(next, 0, PGSIZE);
    if (!cmpxch(entryp, entry, v2p(next) | PTE_P | PTE_W | flags)) {
      kfree((void*) next);
      goto retry;
    }
  }
  return next;
}

// Return the address of the PTE in page table pgdir
// that corresponds to linear address va.  If create!=0,
// create any required page table pages.
atomic<pme_t>*
walkpgdir(pgmap *pml4, u64 va, int create)
{
  auto pdp = descend(pml4, va, PTE_U, create, 3);
  if (pdp == NULL)
    return NULL;
  auto pd = descend(pdp, va, PTE_U, create, 2);
  if (pd == NULL)
    return NULL;
  auto pt = descend(pd, va, PTE_U, create, 1);
  if (pt == NULL)
    return NULL;
  return &pt->e[PX(0,va)];
}

// Map from 0 to 128Gbytes starting at KBASE.
void
initpg(void)
{
  extern char end[];
  u64 va = KBASE;
  paddr pa = 0;

  while (va < (KBASE+(128ull<<30))) {
    auto pdp = descend(&kpml4, va, 0, 1, 3);
    auto pd = descend(pdp, va, 0, 1, 2);
    atomic<pme_t> *sp = &pd->e[PX(1,va)];
    u64 flags = PTE_W | PTE_P | PTE_PS;
    // Set NX for non-code pages
    if (va >= (u64) end)
      flags |= PTE_NX;
    *sp = pa | flags;
    va = va + PGSIZE*512;
    pa += PGSIZE*512;
  }
}

// Set up kernel part of a page table.
pgmap*
setupkvm(void)
{
  pgmap *pml4;
  int k;

  if((pml4 = (pgmap*)kalloc()) == 0)
    return 0;
  k = PX(3, KBASE);
  memset(&pml4->e[0], 0, 8*k);
  memmove(&pml4->e[k], &kpml4.e[k], 8*(512-k));
  return pml4;
}

int
setupkshared(pgmap *pml4, char *kshared)
{
  for (u64 off = 0; off < KSHAREDSIZE; off+=4096) {
    atomic<pme_t> *pte = walkpgdir(pml4, (u64) (KSHARED+off), 1);
    if (pte == NULL)
      panic("setupkshared: oops");
    *pte = v2p(kshared+off) | PTE_P | PTE_U | PTE_W;
  }
  return 0;
}

// Switch h/w page table register to the kernel-only page table,
// for when no process is running.
void
switchkvm(void)
{
  lcr3(v2p(&kpml4));   // switch to the kernel page table
}

// Switch TSS and h/w page table to correspond to process p.
void
switchuvm(struct proc *p)
{
  u64 base = (u64) &mycpu()->ts;
  pushcli();
  mycpu()->gdt[TSSSEG>>3] = (struct segdesc)
    SEGDESC(base, (sizeof(mycpu()->ts)-1), SEG_P|SEG_TSS64A);
  mycpu()->gdt[(TSSSEG>>3)+1] = (struct segdesc) SEGDESCHI(base);
  mycpu()->ts.rsp[0] = (u64) myproc()->kstack + KSTACKSIZE;
  mycpu()->ts.iomba = (u16)offsetof(struct taskstate, iopb);
  ltr(TSSSEG);
  if(p->vmap == 0 || p->vmap->pml4 == 0)
    panic("switchuvm: no vmap/pml4");
  lcr3(v2p(p->vmap->pml4));  // switch to new address space
  popcli();
}

static void
freepm(pgmap *pm, int level)
{
  int i;

  if (level != 0) {
    for (i = 0; i < 512; i++) {
      pme_t entry = pm->e[i];
      if (entry & PTE_P)
        freepm((pgmap*) p2v(PTE_ADDR(entry)), level - 1);
    }
  }

  kfree(pm);
}

// Free a page table and all the physical memory pages
// in the user part.
void
freevm(pgmap *pml4)
{
  int k;
  int i;

  if(pml4 == 0)
    panic("freevm: no pgdir");

  // Don't free kernel portion of the pml4
  k = PX(3, KBASE);
  for (i = 0; i < k; i++) {
    pme_t entry = pml4->e[i];
    if (entry & PTE_P) {
      freepm((pgmap*) p2v(PTE_ADDR(entry)), 2);
    }
  }

  kfree(pml4);
}

// Set up CPU's kernel segment descriptors.
// Run once at boot time on each CPU.
void
inittls(void)
{
  struct cpu *c;

  // Initialize cpu-local storage.
  c = &cpus[cpunum()];
  writegs(KDSEG);
  writemsr(MSR_GS_BASE, (u64)&c->cpu);
  c->cpu = c;
  c->proc = NULL;
  c->kmem = &kmems[cpunum()];
}

atomic<u64> tlbflush_req;

void
tlbflush()
{
  u64 myreq = tlbflush_req++;

  pushcli();
  // the caller may not hold any spinlock, because other CPUs might
  // be spinning waiting for that spinlock, with interrupts disabled,
  // so we will deadlock waiting for their TLB flush..
  assert(mycpu()->ncli == 1);

  int myid = mycpu()->id;
  lcr3(rcr3());
  popcli();

  for (int i = 0; i < ncpu; i++)
    if (i != myid)
      lapic_tlbflush(i);

  for (int i = 0; i < ncpu; i++)
    if (i != myid)
      while (cpus[i].tlbflush_done < myreq)
        /* spin */ ;
}
<commit_msg>avoid pushcli/popcli<commit_after>#include "types.h"
#include "amd64.h"
#include "mmu.h"
#include "cpu.hh"
#include "kernel.hh"
#include "bits.hh"
#include "spinlock.h"
#include "kalloc.hh"
#include "queue.h"
#include "condvar.h"
#include "proc.hh"
#include "vm.hh"
#include <stddef.h>

using namespace std;

static pgmap*
descend(pgmap *dir, u64 va, u64 flags, int create, int level)
{
  atomic<pme_t> *entryp;
  pme_t entry;
  pgmap *next;

retry:
  entryp = &dir->e[PX(level, va)];
  entry = entryp->load();
  if (entry & PTE_P) {
    next = (pgmap*) p2v(PTE_ADDR(entry));
  } else {
    if (!create)
      return NULL;
    next = (pgmap*) kalloc();
    if (!next)
      return NULL;
    memset(next, 0, PGSIZE);
    if (!cmpxch(entryp, entry, v2p(next) | PTE_P | PTE_W | flags)) {
      kfree((void*) next);
      goto retry;
    }
  }
  return next;
}

// Return the address of the PTE in page table pgdir
// that corresponds to linear address va.  If create!=0,
// create any required page table pages.
atomic<pme_t>*
walkpgdir(pgmap *pml4, u64 va, int create)
{
  auto pdp = descend(pml4, va, PTE_U, create, 3);
  if (pdp == NULL)
    return NULL;
  auto pd = descend(pdp, va, PTE_U, create, 2);
  if (pd == NULL)
    return NULL;
  auto pt = descend(pd, va, PTE_U, create, 1);
  if (pt == NULL)
    return NULL;
  return &pt->e[PX(0,va)];
}

// Map from 0 to 128Gbytes starting at KBASE.
void
initpg(void)
{
  extern char end[];
  u64 va = KBASE;
  paddr pa = 0;

  while (va < (KBASE+(128ull<<30))) {
    auto pdp = descend(&kpml4, va, 0, 1, 3);
    auto pd = descend(pdp, va, 0, 1, 2);
    atomic<pme_t> *sp = &pd->e[PX(1,va)];
    u64 flags = PTE_W | PTE_P | PTE_PS;
    // Set NX for non-code pages
    if (va >= (u64) end)
      flags |= PTE_NX;
    *sp = pa | flags;
    va = va + PGSIZE*512;
    pa += PGSIZE*512;
  }
}

// Set up kernel part of a page table.
pgmap*
setupkvm(void)
{
  pgmap *pml4;
  int k;

  if((pml4 = (pgmap*)kalloc()) == 0)
    return 0;
  k = PX(3, KBASE);
  memset(&pml4->e[0], 0, 8*k);
  memmove(&pml4->e[k], &kpml4.e[k], 8*(512-k));
  return pml4;
}

int
setupkshared(pgmap *pml4, char *kshared)
{
  for (u64 off = 0; off < KSHAREDSIZE; off+=4096) {
    atomic<pme_t> *pte = walkpgdir(pml4, (u64) (KSHARED+off), 1);
    if (pte == NULL)
      panic("setupkshared: oops");
    *pte = v2p(kshared+off) | PTE_P | PTE_U | PTE_W;
  }
  return 0;
}

// Switch h/w page table register to the kernel-only page table,
// for when no process is running.
void
switchkvm(void)
{
  lcr3(v2p(&kpml4));   // switch to the kernel page table
}

// Switch TSS and h/w page table to correspond to process p.
void
switchuvm(struct proc *p)
{
  u64 base = (u64) &mycpu()->ts;
  pushcli();
  mycpu()->gdt[TSSSEG>>3] = (struct segdesc)
    SEGDESC(base, (sizeof(mycpu()->ts)-1), SEG_P|SEG_TSS64A);
  mycpu()->gdt[(TSSSEG>>3)+1] = (struct segdesc) SEGDESCHI(base);
  mycpu()->ts.rsp[0] = (u64) myproc()->kstack + KSTACKSIZE;
  mycpu()->ts.iomba = (u16)offsetof(struct taskstate, iopb);
  ltr(TSSSEG);
  if(p->vmap == 0 || p->vmap->pml4 == 0)
    panic("switchuvm: no vmap/pml4");
  lcr3(v2p(p->vmap->pml4));  // switch to new address space
  popcli();
}

static void
freepm(pgmap *pm, int level)
{
  int i;

  if (level != 0) {
    for (i = 0; i < 512; i++) {
      pme_t entry = pm->e[i];
      if (entry & PTE_P)
        freepm((pgmap*) p2v(PTE_ADDR(entry)), level - 1);
    }
  }

  kfree(pm);
}

// Free a page table and all the physical memory pages
// in the user part.
void
freevm(pgmap *pml4)
{
  int k;
  int i;

  if(pml4 == 0)
    panic("freevm: no pgdir");

  // Don't free kernel portion of the pml4
  k = PX(3, KBASE);
  for (i = 0; i < k; i++) {
    pme_t entry = pml4->e[i];
    if (entry & PTE_P) {
      freepm((pgmap*) p2v(PTE_ADDR(entry)), 2);
    }
  }

  kfree(pml4);
}

// Set up CPU's kernel segment descriptors.
// Run once at boot time on each CPU.
void
inittls(void)
{
  struct cpu *c;

  // Initialize cpu-local storage.
  c = &cpus[cpunum()];
  writegs(KDSEG);
  writemsr(MSR_GS_BASE, (u64)&c->cpu);
  c->cpu = c;
  c->proc = NULL;
  c->kmem = &kmems[cpunum()];
}

atomic<u64> tlbflush_req;

void
tlbflush()
{
  u64 myreq = tlbflush_req++;

  // the caller may not hold any spinlock, because other CPUs might
  // be spinning waiting for that spinlock, with interrupts disabled,
  // so we will deadlock waiting for their TLB flush..
  assert(mycpu()->ncli == 0);

  for (int i = 0; i < ncpu; i++)
    lapic_tlbflush(i);

  for (int i = 0; i < ncpu; i++)
    while (cpus[i].tlbflush_done < myreq)
      /* spin */ ;
}
<|endoftext|>


<commit_before>
#include <iostream>
#include <string>

#include <thread>
#include <mutex>

#include <boost/program_options.hpp>

#include "filter/compiler.h"

#include "fileemitor.h"
#include "version.h"
#include "worker.h"
namespace po = boost::program_options;

std::mutex screenMutex;

std::string recordSeparator = "\n";
void outDocument(const std::string &what) {
    std::lock_guard<std::mutex> screenLock(screenMutex);
    std::cout.write(what.data(), what.size());
    std::cout << recordSeparator;
}

void updateSeparator(std::string &sep) {
    std::string res;
    bool esc = false;
    for(char c : sep) {
        if (esc) {
            switch(c) {
                case 'n':
                    res += '\n';
                    break;
                case 't':
                    res += '\t';
                    break;
                default:
                    res += c;
            }
            esc = false;
            continue;
        }
        if (c == '\\') {
            esc = true;
        } else {
            res += c;
        }
    }
    sep = res;
}

void correctJobsNumber(u_int &jobs) {
    if (jobs < 1) {
        std::cerr << "hint: adjusting threads number to 1" << std::endl;
        jobs = 1;
    } else if (jobs > 10) {
        std::cerr << "hint: do not be so greedy. adjusting threads number to 10" << std::endl;
        jobs = 10;
    }
}

int main(int argc, const char * argv[]) {

    std::cout.sync_with_stdio(false);

    std::string condition;
    int limit = -1;
    u_int jobs = 1;
    std::string fields;
    bool printProcessingFile = false;
    bool countMode = false;
    bool disableParseLoop = false;
    bool displayVersionAndExit = false;
    std::string fieldSeparator = "\t";

    po::options_description desc("Allowed options");
    desc.add_options()
        ("input-file,f", po::value< std::vector<std::string> >(), "Input files")
        ("condition,c", po::value< std::string >(&condition), "Expression")
        ("limit,n", po::value< int >(&limit)->default_value(-1), "Maximum number of records (default -1 means no limit)")
        ("fields,l", po::value< std::string >(&fields), "Fields to output")
        ("print-file", po::bool_switch(&printProcessingFile), "Print name of processing file")
        ("jobs,j", po::value< u_int >(&jobs)->default_value(1), "Number of threads to run")
        ("count-only", po::bool_switch(&countMode), "Count of matched records, don't print them")
        ("record-separator", po::value<std::string>(&recordSeparator)->default_value("\\n"), "Record separator (\\n by default)")
        ("field-separator", po::value<std::string>(&fieldSeparator)->default_value("\\t"), "Field separator for TSV output (\\t by default)")
        ("disable-parse-loop", po::bool_switch(&disableParseLoop), "Disable experimental parsing mode (enabled by default)")
        ("help,h", "Show help message")
        ("version", po::bool_switch(&displayVersionAndExit), "Display version and exit")
    ;

    po::positional_options_description p;
    p.add("input-file", -1);

    po::variables_map vm;
    try {
        po::store(po::command_line_parser(argc, argv).
                  options(desc).positional(p).run(), vm);
        po::notify(vm);
    } catch(const boost::program_options::error &e) {
        std::cerr << "Sorry, I couldn't parse arguments: " << e.what() << std::endl;
        std::cerr << "Try --help" << std::endl;
        return 1;
    }
    if (vm.count("help")) {
        std::cout << desc << std::endl;
        return 0;
    }
    if (displayVersionAndExit) {
        std::cout << "aq version: " << avroqVersion() << std::endl;
        return 0;
    }

    updateSeparator(recordSeparator);
    updateSeparator(fieldSeparator);

    filter::Compiler filterCompiler;
    std::shared_ptr<filter::Filter> filter;

    if (!condition.empty()) {
        try {
            filter = filterCompiler.compile(condition);
        } catch(const filter::Compiler::CompileError &e) {
            std::cerr << "Condition complile failed: " << std::endl;
            std::cerr << e.what() << std::endl;
            return 1;
        }
    }

    if (vm.count("input-file")) {

        const auto &fileList = vm["input-file"].as< std::vector<std::string> >();

        FileEmitor emitor(fileList, limit, outDocument);

        emitor.setFilter(filter);
        emitor.setTsvFieldList(fields, fieldSeparator);
        if (printProcessingFile) {
            emitor.enablePrintProcessingFile();
        }
        if (countMode) {
            emitor.enableCountOnlyMode();
        }
        if ( ! disableParseLoop ) {
            emitor.enableParseLoop();
        }
        std::vector<std::thread> workers;

        correctJobsNumber(jobs);

        for(u_int i = 0; i < jobs; ++i) { // TODO: check for inadequate values
            workers.emplace_back(
                    std::thread(Worker(emitor))
                );
        }

        for(auto &p : workers) {
            p.join();
        }

        if (countMode) {
            std::cout << "Matched documents: " << emitor.getCountedDocuments() << std::endl;
        }

        if (emitor.getLastError().size() > 0) {
            std::cerr << emitor.getLastError() << std::endl;
            return 1;
        }

    } else {
        std::cerr << "No input files" << std::endl;
        return 1;
    }

    return 0;
}
<commit_msg>Increased threads limit to 16, for 16-x core systems<commit_after>
#include <iostream>
#include <string>

#include <thread>
#include <mutex>

#include <boost/program_options.hpp>

#include "filter/compiler.h"

#include "fileemitor.h"
#include "version.h"
#include "worker.h"
namespace po = boost::program_options;

std::mutex screenMutex;

std::string recordSeparator = "\n";
void outDocument(const std::string &what) {
    std::lock_guard<std::mutex> screenLock(screenMutex);
    std::cout.write(what.data(), what.size());
    std::cout << recordSeparator;
}

void updateSeparator(std::string &sep) {
    std::string res;
    bool esc = false;
    for(char c : sep) {
        if (esc) {
            switch(c) {
                case 'n':
                    res += '\n';
                    break;
                case 't':
                    res += '\t';
                    break;
                default:
                    res += c;
            }
            esc = false;
            continue;
        }
        if (c == '\\') {
            esc = true;
        } else {
            res += c;
        }
    }
    sep = res;
}

void correctJobsNumber(u_int &jobs) {
    if (jobs < 1) {
        std::cerr << "hint: adjusting threads number to 1" << std::endl;
        jobs = 1;
    } else if (jobs > 16) {
        std::cerr << "hint: do not be so greedy. adjusting threads number to 16" << std::endl;
        jobs = 16;
    }
}

int main(int argc, const char * argv[]) {

    std::cout.sync_with_stdio(false);

    std::string condition;
    int limit = -1;
    u_int jobs = 1;
    std::string fields;
    bool printProcessingFile = false;
    bool countMode = false;
    bool disableParseLoop = false;
    bool displayVersionAndExit = false;
    std::string fieldSeparator = "\t";

    po::options_description desc("Allowed options");
    desc.add_options()
        ("input-file,f", po::value< std::vector<std::string> >(), "Input files")
        ("condition,c", po::value< std::string >(&condition), "Expression")
        ("limit,n", po::value< int >(&limit)->default_value(-1), "Maximum number of records (default -1 means no limit)")
        ("fields,l", po::value< std::string >(&fields), "Fields to output")
        ("print-file", po::bool_switch(&printProcessingFile), "Print name of processing file")
        ("jobs,j", po::value< u_int >(&jobs)->default_value(1), "Number of threads to run")
        ("count-only", po::bool_switch(&countMode), "Count of matched records, don't print them")
        ("record-separator", po::value<std::string>(&recordSeparator)->default_value("\\n"), "Record separator (\\n by default)")
        ("field-separator", po::value<std::string>(&fieldSeparator)->default_value("\\t"), "Field separator for TSV output (\\t by default)")
        ("disable-parse-loop", po::bool_switch(&disableParseLoop), "Disable experimental parsing mode (enabled by default)")
        ("help,h", "Show help message")
        ("version", po::bool_switch(&displayVersionAndExit), "Display version and exit")
    ;

    po::positional_options_description p;
    p.add("input-file", -1);

    po::variables_map vm;
    try {
        po::store(po::command_line_parser(argc, argv).
                  options(desc).positional(p).run(), vm);
        po::notify(vm);
    } catch(const boost::program_options::error &e) {
        std::cerr << "Sorry, I couldn't parse arguments: " << e.what() << std::endl;
        std::cerr << "Try --help" << std::endl;
        return 1;
    }
    if (vm.count("help")) {
        std::cout << desc << std::endl;
        return 0;
    }
    if (displayVersionAndExit) {
        std::cout << "aq version: " << avroqVersion() << std::endl;
        return 0;
    }

    updateSeparator(recordSeparator);
    updateSeparator(fieldSeparator);

    filter::Compiler filterCompiler;
    std::shared_ptr<filter::Filter> filter;

    if (!condition.empty()) {
        try {
            filter = filterCompiler.compile(condition);
        } catch(const filter::Compiler::CompileError &e) {
            std::cerr << "Condition complile failed: " << std::endl;
            std::cerr << e.what() << std::endl;
            return 1;
        }
    }

    if (vm.count("input-file")) {

        const auto &fileList = vm["input-file"].as< std::vector<std::string> >();

        FileEmitor emitor(fileList, limit, outDocument);

        emitor.setFilter(filter);
        emitor.setTsvFieldList(fields, fieldSeparator);
        if (printProcessingFile) {
            emitor.enablePrintProcessingFile();
        }
        if (countMode) {
            emitor.enableCountOnlyMode();
        }
        if ( ! disableParseLoop ) {
            emitor.enableParseLoop();
        }
        std::vector<std::thread> workers;

        correctJobsNumber(jobs);

        for(u_int i = 0; i < jobs; ++i) { // TODO: check for inadequate values
            workers.emplace_back(
                    std::thread(Worker(emitor))
                );
        }

        for(auto &p : workers) {
            p.join();
        }

        if (countMode) {
            std::cout << "Matched documents: " << emitor.getCountedDocuments() << std::endl;
        }

        if (emitor.getLastError().size() > 0) {
            std::cerr << emitor.getLastError() << std::endl;
            return 1;
        }

    } else {
        std::cerr << "No input files" << std::endl;
        return 1;
    }

    return 0;
}
<|endoftext|>


<commit_before>// Copyright 2016 The TensorFlow Authors. All Rights Reserved.
//
// 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.
// =============================================================================
// FinishedNodes returns a 1-D tensor listing the nodes that are finished
// accumulating.
#include "tensorflow/contrib/tensor_forest/core/ops/tree_utils.h"

#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/shape_inference.h"
#include "tensorflow/core/kernels/bounds_check.h"
#include "tensorflow/core/lib/random/simple_philox.h"
#include "tensorflow/core/platform/env.h"
#include "tensorflow/core/platform/mutex.h"
#include "tensorflow/core/util/work_sharder.h"

namespace tensorflow {

using shape_inference::Dimension;
using shape_inference::InferenceContext;
using shape_inference::Shape;

using std::placeholders::_1;
using std::placeholders::_2;

using tensorforest::CheckTensorBounds;
using tensorforest::Sum;
using tensorforest::BestSplitDominatesClassificationBootstrap;
using tensorforest::BestSplitDominatesClassificationChebyshev;
using tensorforest::BestSplitDominatesClassificationHoeffding;
using tensorforest::BestSplitDominatesRegression;

namespace {

struct EvaluateParams {
  Tensor leaves;
  Tensor node_to_accumulator;
  Tensor accumulator_sums;
  Tensor birth_epochs;
  int current_epoch;
  int32 num_split_after_samples;
  int32 min_split_samples;
  bool need_random;
  int64 random_seed;
  std::function<bool(int, random::SimplePhilox*)> dominate_method;
};

void Evaluate(const EvaluateParams& params, mutex* mutex, int32 start,
              int32 end, std::vector<int32>* final_finished_leaves,
              std::vector<int32>* final_stale) {
  const auto leaves = params.leaves.unaligned_flat<int32>();
  const auto node_map = params.node_to_accumulator.unaligned_flat<int32>();
  const auto sums = params.accumulator_sums.tensor<float, 2>();
  const auto start_epochs = params.birth_epochs.unaligned_flat<int32>();

  const int32 num_accumulators =
      static_cast<int32>(params.accumulator_sums.shape().dim_size(0));

  std::vector<int32> finished_leaves;
  std::vector<int32> stale;

  std::unique_ptr<random::SimplePhilox> simple_philox;
  random::PhiloxRandom rnd_gen(params.random_seed);

  if (params.need_random) {
    simple_philox.reset(new random::SimplePhilox(&rnd_gen));
  }

  for (int32 i = start; i < end; i++) {
    const int32 leaf = internal::SubtleMustCopy(leaves(i));
    if (leaf == -1) {
      continue;
    }
    if (!FastBoundsCheck(leaf, node_map.size())) {
      LOG(ERROR) << "leaf " << leaf << " not in valid range.";
    }
    const int32 accumulator = internal::SubtleMustCopy(node_map(leaf));
    if (accumulator < 0) {
      continue;
    }

    if (!FastBoundsCheck(accumulator, num_accumulators)) {
      LOG(ERROR) << "accumulator " << accumulator << " not in valid range.";
    }
    // The first column holds the number of samples seen.
    // For classification, this should be the sum of the other columns.
    int32 count = sums(accumulator, 0);

    if (params.current_epoch > start_epochs(leaf) + 1) {
      if (count >= params.min_split_samples) {
        finished_leaves.push_back(leaf);
      } else {
        stale.push_back(leaf);
      }
      continue;
    }

    if (count >= params.num_split_after_samples) {
      finished_leaves.push_back(leaf);
      continue;
    }

    if (count < params.min_split_samples) {
      continue;
    }

    bool finished = params.dominate_method(accumulator, simple_philox.get());
    if (finished) {
      finished_leaves.push_back(leaf);
    }
  }
  mutex_lock m(*mutex);
  final_finished_leaves->insert(final_finished_leaves->end(),
                                finished_leaves.begin(), finished_leaves.end());
  final_stale->insert(final_stale->end(), stale.begin(), stale.end());
}
}  // namespace

REGISTER_OP("FinishedNodes")
    .Attr("regression: bool = false")
    .Attr("num_split_after_samples: int")
    .Attr("min_split_samples: int")
    .Attr("dominate_fraction: float = 0.99")
    // TODO(thomaswc): Test out bootstrap on several datasets, confirm it
    // works well, make it the default.
    .Attr(
        "dominate_method:"
        " {'none', 'hoeffding', 'bootstrap', 'chebyshev'} = 'hoeffding'")
    .Attr("random_seed: int = 0")
    .Input("leaves: int32")
    .Input("node_to_accumulator: int32")
    .Input("split_sums: float")
    .Input("split_squares: float")
    .Input("accumulator_sums: float")
    .Input("accumulator_squares: float")
    .Input("birth_epochs: int32")
    .Input("current_epoch: int32")
    .Output("finished: int32")
    .Output("stale: int32")
    .SetShapeFn([](InferenceContext* c) {
      c->set_output(0, c->Vector(InferenceContext::kUnknownDim));
      c->set_output(1, c->Vector(InferenceContext::kUnknownDim));
      return Status::OK();
    })
    .Doc(R"doc(
Determines which of the given leaf nodes are done accumulating.

leaves:= A 1-d int32 tensor.  Lists the nodes that are currently leaves.
node_to_accumulator: If the i-th node is fertile, `node_to_accumulator[i]`
  is it's accumulator slot.  Otherwise, `node_to_accumulator[i]` is -1.
split_sums:= a 3-d tensor where `split_sums[a][s]` summarizes the
  training labels for examples that fall into the fertile node associated with
  accumulator slot s and have then taken the *left* branch of candidate split
  s.  For a classification problem, `split_sums[a][s][c]` is the count of such
  examples with class c and for regression problems, `split_sums[a][s]` is the
  sum of the regression labels for such examples.
split_squares: Same as split_sums, but it contains the sum of the
  squares of the regression labels.  Only used for regression.  For
  classification problems, pass a dummy tensor into this.
accumulator_sums: For classification, `accumulator_sums[a][c]` records how
  many training examples have class c and have ended up in the fertile node
  associated with accumulator slot a.  It has the total sum in entry 0 for
  convenience. For regression, it is the same except it contains the sum
  of the input labels that have been seen, and entry 0 contains the number
  of training examples that have been seen.
accumulator_squares: Same as accumulator_sums, but it contains the sum of the
  squares of the regression labels.  Only used for regression.  For
  classification problems, pass a dummy tensor into this.
birth_epochs:= A 1-d int32 tensor.  `birth_epochs[i]` contains the epoch
  the i-th node was created in.
current_epoch:= A 1-d int32 tensor with shape (1).  `current_epoch[0]`
  stores the current epoch number.
finished:= A 1-d int32 tensor containing the indices of the finished nodes.
  Nodes are finished if they have received at least num_split_after_samples
  samples, or if they have received min_split_samples and the best scoring
  split is sufficiently greater than the next best split.
stale:= A 1-d int32 tensor containing the fertile nodes that were created two
  or more epochs ago.

)doc");

class FinishedNodes : public OpKernel {
 public:
  explicit FinishedNodes(OpKernelConstruction* context)
      : OpKernel(context) {
    OP_REQUIRES_OK(context, context->GetAttr(
        "regression", &regression_));
    OP_REQUIRES_OK(context, context->GetAttr(
        "num_split_after_samples", &num_split_after_samples_));
    OP_REQUIRES_OK(context, context->GetAttr(
        "min_split_samples", &min_split_samples_));
    OP_REQUIRES_OK(context, context->GetAttr(
        "dominate_fraction", &dominate_fraction_));
    OP_REQUIRES_OK(context,
                   context->GetAttr("dominate_method", &dominate_method_));
    OP_REQUIRES_OK(context, context->GetAttr("random_seed", &random_seed_));
  }

  void Compute(OpKernelContext* context) override {
    const Tensor& leaf_tensor = context->input(0);
    const Tensor& node_to_accumulator = context->input(1);
    const Tensor& split_sums = context->input(2);
    const Tensor& split_squares = context->input(3);
    const Tensor& accumulator_sums = context->input(4);
    const Tensor& accumulator_squares = context->input(5);
    const Tensor& birth_epochs = context->input(6);
    const Tensor& current_epoch = context->input(7);

    OP_REQUIRES(context, leaf_tensor.shape().dims() == 1,
                errors::InvalidArgument(
                    "leaf_tensor should be one-dimensional"));
    OP_REQUIRES(context, node_to_accumulator.shape().dims() == 1,
                errors::InvalidArgument(
                    "node_to_accumulator should be one-dimensional"));
    OP_REQUIRES(context, split_sums.shape().dims() == 3,
                errors::InvalidArgument(
                    "split_sums should be three-dimensional"));
    OP_REQUIRES(context, accumulator_sums.shape().dims() == 2,
                errors::InvalidArgument(
                    "accumulator_sums should be two-dimensional"));
    OP_REQUIRES(context, birth_epochs.shape().dims() == 1,
                errors::InvalidArgument(
                    "birth_epochs should be one-dimensional"));
    OP_REQUIRES(
        context,
        birth_epochs.shape().dim_size(0) ==
        node_to_accumulator.shape().dim_size(0),
        errors::InvalidArgument(
            "birth_epochs and node_to_accumulator should be the same size."));

    // Check tensor bounds.
    if (!CheckTensorBounds(context, leaf_tensor)) return;
    if (!CheckTensorBounds(context, node_to_accumulator)) return;
    if (!CheckTensorBounds(context, split_sums)) return;
    if (!CheckTensorBounds(context, split_squares)) return;
    if (!CheckTensorBounds(context, accumulator_sums)) return;
    if (!CheckTensorBounds(context, accumulator_squares)) return;
    if (!CheckTensorBounds(context, birth_epochs)) return;
    if (!CheckTensorBounds(context, current_epoch)) return;

    const int32 epoch = current_epoch.unaligned_flat<int32>()(0);

    const int32 num_leaves = static_cast<int32>(
        leaf_tensor.shape().dim_size(0));

    auto worker_threads = context->device()->tensorflow_cpu_worker_threads();
    int num_threads = worker_threads->num_threads;

    EvaluateParams params;
    params.leaves = leaf_tensor;
    params.node_to_accumulator = node_to_accumulator;
    params.accumulator_sums = accumulator_sums;
    params.birth_epochs = birth_epochs;
    params.current_epoch = epoch;
    params.min_split_samples = min_split_samples_;
    params.num_split_after_samples = num_split_after_samples_;
    params.need_random = false;

    if (regression_) {
      params.dominate_method =
          std::bind(&BestSplitDominatesRegression, accumulator_sums,
                    accumulator_squares, split_sums, split_squares, _1);
    } else {
      if (dominate_method_ == "none") {
        params.dominate_method = [](int, random::SimplePhilox*) {
          return false;
        };
      } else if (dominate_method_ == "hoeffding") {
        params.dominate_method =
            std::bind(&BestSplitDominatesClassificationHoeffding,
                      accumulator_sums, split_sums, _1, dominate_fraction_);
      } else if (dominate_method_ == "chebyshev") {
        params.dominate_method =
            std::bind(&BestSplitDominatesClassificationChebyshev,
                      accumulator_sums, split_sums, _1, dominate_fraction_);
      } else if (dominate_method_ == "bootstrap") {
        params.need_random = true;

        params.random_seed = random_seed_;
        if (params.random_seed == 0) {
          params.random_seed = static_cast<uint64>(Env::Default()->NowMicros());
        }

        params.dominate_method =
            std::bind(&BestSplitDominatesClassificationBootstrap,
                      accumulator_sums, split_sums, _1, dominate_fraction_, _2);
      } else {
        LOG(FATAL) << "Unknown dominate method " << dominate_method_;
      }
    }

    std::vector<int32> finished_leaves;
    std::vector<int32> stale;
    mutex m;
    // Require at least 100 leaves per thread.  I guess that's about 800 cost
    // per unit.  This isn't well defined.
    const int64 costPerUnit = 800;
    auto work = [&params, &finished_leaves, &stale, &m, num_leaves](int64 start,
                                                                    int64 end) {
      CHECK(start <= end);
      CHECK(end <= num_leaves);
      Evaluate(params, &m, static_cast<int32>(start), static_cast<int32>(end),
               &finished_leaves, &stale);
    };
    Shard(num_threads, worker_threads->workers, num_leaves, costPerUnit, work);

    // Copy to output.
    Tensor* output_finished = nullptr;
    TensorShape finished_shape;
    finished_shape.AddDim(finished_leaves.size());
    OP_REQUIRES_OK(context,
                   context->allocate_output(0, finished_shape,
                                            &output_finished));
    auto out_finished = output_finished->unaligned_flat<int32>();
    std::copy(finished_leaves.begin(), finished_leaves.end(),
              out_finished.data());

    Tensor* output_stale = nullptr;
    TensorShape stale_shape;
    stale_shape.AddDim(stale.size());
    OP_REQUIRES_OK(context,
                   context->allocate_output(1, stale_shape,
                                            &output_stale));
    auto out_stale = output_stale->unaligned_flat<int32>();
    std::copy(stale.begin(), stale.end(), out_stale.data());
  }

 private:
  bool regression_;
  int32 num_split_after_samples_;
  int32 min_split_samples_;
  float dominate_fraction_;
  string dominate_method_;
  int32 random_seed_;
};

REGISTER_KERNEL_BUILDER(Name("FinishedNodes").Device(DEVICE_CPU),
                        FinishedNodes);

}  // namespace tensorflow
<commit_msg>Change default dominate_method to bootsrap. Change: 137959663<commit_after>// Copyright 2016 The TensorFlow Authors. All Rights Reserved.
//
// 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.
// =============================================================================
// FinishedNodes returns a 1-D tensor listing the nodes that are finished
// accumulating.
#include "tensorflow/contrib/tensor_forest/core/ops/tree_utils.h"

#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/shape_inference.h"
#include "tensorflow/core/kernels/bounds_check.h"
#include "tensorflow/core/lib/random/simple_philox.h"
#include "tensorflow/core/platform/env.h"
#include "tensorflow/core/platform/mutex.h"
#include "tensorflow/core/util/work_sharder.h"

namespace tensorflow {

using shape_inference::Dimension;
using shape_inference::InferenceContext;
using shape_inference::Shape;

using std::placeholders::_1;
using std::placeholders::_2;

using tensorforest::CheckTensorBounds;
using tensorforest::Sum;
using tensorforest::BestSplitDominatesClassificationBootstrap;
using tensorforest::BestSplitDominatesClassificationChebyshev;
using tensorforest::BestSplitDominatesClassificationHoeffding;
using tensorforest::BestSplitDominatesRegression;

namespace {

struct EvaluateParams {
  Tensor leaves;
  Tensor node_to_accumulator;
  Tensor accumulator_sums;
  Tensor birth_epochs;
  int current_epoch;
  int32 num_split_after_samples;
  int32 min_split_samples;
  bool need_random;
  int64 random_seed;
  std::function<bool(int, random::SimplePhilox*)> dominate_method;
};

void Evaluate(const EvaluateParams& params, mutex* mutex, int32 start,
              int32 end, std::vector<int32>* final_finished_leaves,
              std::vector<int32>* final_stale) {
  const auto leaves = params.leaves.unaligned_flat<int32>();
  const auto node_map = params.node_to_accumulator.unaligned_flat<int32>();
  const auto sums = params.accumulator_sums.tensor<float, 2>();
  const auto start_epochs = params.birth_epochs.unaligned_flat<int32>();

  const int32 num_accumulators =
      static_cast<int32>(params.accumulator_sums.shape().dim_size(0));

  std::vector<int32> finished_leaves;
  std::vector<int32> stale;

  std::unique_ptr<random::SimplePhilox> simple_philox;
  random::PhiloxRandom rnd_gen(params.random_seed);

  if (params.need_random) {
    simple_philox.reset(new random::SimplePhilox(&rnd_gen));
  }

  for (int32 i = start; i < end; i++) {
    const int32 leaf = internal::SubtleMustCopy(leaves(i));
    if (leaf == -1) {
      continue;
    }
    if (!FastBoundsCheck(leaf, node_map.size())) {
      LOG(ERROR) << "leaf " << leaf << " not in valid range.";
    }
    const int32 accumulator = internal::SubtleMustCopy(node_map(leaf));
    if (accumulator < 0) {
      continue;
    }

    if (!FastBoundsCheck(accumulator, num_accumulators)) {
      LOG(ERROR) << "accumulator " << accumulator << " not in valid range.";
    }
    // The first column holds the number of samples seen.
    // For classification, this should be the sum of the other columns.
    int32 count = sums(accumulator, 0);

    if (params.current_epoch > start_epochs(leaf) + 1) {
      if (count >= params.min_split_samples) {
        finished_leaves.push_back(leaf);
      } else {
        stale.push_back(leaf);
      }
      continue;
    }

    if (count >= params.num_split_after_samples) {
      finished_leaves.push_back(leaf);
      continue;
    }

    if (count < params.min_split_samples) {
      continue;
    }

    bool finished = params.dominate_method(accumulator, simple_philox.get());
    if (finished) {
      finished_leaves.push_back(leaf);
    }
  }
  mutex_lock m(*mutex);
  final_finished_leaves->insert(final_finished_leaves->end(),
                                finished_leaves.begin(), finished_leaves.end());
  final_stale->insert(final_stale->end(), stale.begin(), stale.end());
}
}  // namespace

REGISTER_OP("FinishedNodes")
    .Attr("regression: bool = false")
    .Attr("num_split_after_samples: int")
    .Attr("min_split_samples: int")
    .Attr("dominate_fraction: float = 0.99")
    .Attr(
        "dominate_method:"
        " {'none', 'hoeffding', 'bootstrap', 'chebyshev'} = 'bootstrap'")
    .Attr("random_seed: int = 0")
    .Input("leaves: int32")
    .Input("node_to_accumulator: int32")
    .Input("split_sums: float")
    .Input("split_squares: float")
    .Input("accumulator_sums: float")
    .Input("accumulator_squares: float")
    .Input("birth_epochs: int32")
    .Input("current_epoch: int32")
    .Output("finished: int32")
    .Output("stale: int32")
    .SetShapeFn([](InferenceContext* c) {
      c->set_output(0, c->Vector(InferenceContext::kUnknownDim));
      c->set_output(1, c->Vector(InferenceContext::kUnknownDim));
      return Status::OK();
    })
    .Doc(R"doc(
Determines which of the given leaf nodes are done accumulating.

leaves:= A 1-d int32 tensor.  Lists the nodes that are currently leaves.
node_to_accumulator: If the i-th node is fertile, `node_to_accumulator[i]`
  is it's accumulator slot.  Otherwise, `node_to_accumulator[i]` is -1.
split_sums:= a 3-d tensor where `split_sums[a][s]` summarizes the
  training labels for examples that fall into the fertile node associated with
  accumulator slot s and have then taken the *left* branch of candidate split
  s.  For a classification problem, `split_sums[a][s][c]` is the count of such
  examples with class c and for regression problems, `split_sums[a][s]` is the
  sum of the regression labels for such examples.
split_squares: Same as split_sums, but it contains the sum of the
  squares of the regression labels.  Only used for regression.  For
  classification problems, pass a dummy tensor into this.
accumulator_sums: For classification, `accumulator_sums[a][c]` records how
  many training examples have class c and have ended up in the fertile node
  associated with accumulator slot a.  It has the total sum in entry 0 for
  convenience. For regression, it is the same except it contains the sum
  of the input labels that have been seen, and entry 0 contains the number
  of training examples that have been seen.
accumulator_squares: Same as accumulator_sums, but it contains the sum of the
  squares of the regression labels.  Only used for regression.  For
  classification problems, pass a dummy tensor into this.
birth_epochs:= A 1-d int32 tensor.  `birth_epochs[i]` contains the epoch
  the i-th node was created in.
current_epoch:= A 1-d int32 tensor with shape (1).  `current_epoch[0]`
  stores the current epoch number.
finished:= A 1-d int32 tensor containing the indices of the finished nodes.
  Nodes are finished if they have received at least num_split_after_samples
  samples, or if they have received min_split_samples and the best scoring
  split is sufficiently greater than the next best split.
stale:= A 1-d int32 tensor containing the fertile nodes that were created two
  or more epochs ago.

)doc");

class FinishedNodes : public OpKernel {
 public:
  explicit FinishedNodes(OpKernelConstruction* context)
      : OpKernel(context) {
    OP_REQUIRES_OK(context, context->GetAttr(
        "regression", &regression_));
    OP_REQUIRES_OK(context, context->GetAttr(
        "num_split_after_samples", &num_split_after_samples_));
    OP_REQUIRES_OK(context, context->GetAttr(
        "min_split_samples", &min_split_samples_));
    OP_REQUIRES_OK(context, context->GetAttr(
        "dominate_fraction", &dominate_fraction_));
    OP_REQUIRES_OK(context,
                   context->GetAttr("dominate_method", &dominate_method_));
    OP_REQUIRES_OK(context, context->GetAttr("random_seed", &random_seed_));
  }

  void Compute(OpKernelContext* context) override {
    const Tensor& leaf_tensor = context->input(0);
    const Tensor& node_to_accumulator = context->input(1);
    const Tensor& split_sums = context->input(2);
    const Tensor& split_squares = context->input(3);
    const Tensor& accumulator_sums = context->input(4);
    const Tensor& accumulator_squares = context->input(5);
    const Tensor& birth_epochs = context->input(6);
    const Tensor& current_epoch = context->input(7);

    OP_REQUIRES(context, leaf_tensor.shape().dims() == 1,
                errors::InvalidArgument(
                    "leaf_tensor should be one-dimensional"));
    OP_REQUIRES(context, node_to_accumulator.shape().dims() == 1,
                errors::InvalidArgument(
                    "node_to_accumulator should be one-dimensional"));
    OP_REQUIRES(context, split_sums.shape().dims() == 3,
                errors::InvalidArgument(
                    "split_sums should be three-dimensional"));
    OP_REQUIRES(context, accumulator_sums.shape().dims() == 2,
                errors::InvalidArgument(
                    "accumulator_sums should be two-dimensional"));
    OP_REQUIRES(context, birth_epochs.shape().dims() == 1,
                errors::InvalidArgument(
                    "birth_epochs should be one-dimensional"));
    OP_REQUIRES(
        context,
        birth_epochs.shape().dim_size(0) ==
        node_to_accumulator.shape().dim_size(0),
        errors::InvalidArgument(
            "birth_epochs and node_to_accumulator should be the same size."));

    // Check tensor bounds.
    if (!CheckTensorBounds(context, leaf_tensor)) return;
    if (!CheckTensorBounds(context, node_to_accumulator)) return;
    if (!CheckTensorBounds(context, split_sums)) return;
    if (!CheckTensorBounds(context, split_squares)) return;
    if (!CheckTensorBounds(context, accumulator_sums)) return;
    if (!CheckTensorBounds(context, accumulator_squares)) return;
    if (!CheckTensorBounds(context, birth_epochs)) return;
    if (!CheckTensorBounds(context, current_epoch)) return;

    const int32 epoch = current_epoch.unaligned_flat<int32>()(0);

    const int32 num_leaves = static_cast<int32>(
        leaf_tensor.shape().dim_size(0));

    auto worker_threads = context->device()->tensorflow_cpu_worker_threads();
    int num_threads = worker_threads->num_threads;

    EvaluateParams params;
    params.leaves = leaf_tensor;
    params.node_to_accumulator = node_to_accumulator;
    params.accumulator_sums = accumulator_sums;
    params.birth_epochs = birth_epochs;
    params.current_epoch = epoch;
    params.min_split_samples = min_split_samples_;
    params.num_split_after_samples = num_split_after_samples_;
    params.need_random = false;

    if (regression_) {
      params.dominate_method =
          std::bind(&BestSplitDominatesRegression, accumulator_sums,
                    accumulator_squares, split_sums, split_squares, _1);
    } else {
      if (dominate_method_ == "none") {
        params.dominate_method = [](int, random::SimplePhilox*) {
          return false;
        };
      } else if (dominate_method_ == "hoeffding") {
        params.dominate_method =
            std::bind(&BestSplitDominatesClassificationHoeffding,
                      accumulator_sums, split_sums, _1, dominate_fraction_);
      } else if (dominate_method_ == "chebyshev") {
        params.dominate_method =
            std::bind(&BestSplitDominatesClassificationChebyshev,
                      accumulator_sums, split_sums, _1, dominate_fraction_);
      } else if (dominate_method_ == "bootstrap") {
        params.need_random = true;

        params.random_seed = random_seed_;
        if (params.random_seed == 0) {
          params.random_seed = static_cast<uint64>(Env::Default()->NowMicros());
        }

        params.dominate_method =
            std::bind(&BestSplitDominatesClassificationBootstrap,
                      accumulator_sums, split_sums, _1, dominate_fraction_, _2);
      } else {
        LOG(FATAL) << "Unknown dominate method " << dominate_method_;
      }
    }

    std::vector<int32> finished_leaves;
    std::vector<int32> stale;
    mutex m;
    // Require at least 100 leaves per thread.  I guess that's about 800 cost
    // per unit.  This isn't well defined.
    const int64 costPerUnit = 800;
    auto work = [&params, &finished_leaves, &stale, &m, num_leaves](int64 start,
                                                                    int64 end) {
      CHECK(start <= end);
      CHECK(end <= num_leaves);
      Evaluate(params, &m, static_cast<int32>(start), static_cast<int32>(end),
               &finished_leaves, &stale);
    };
    Shard(num_threads, worker_threads->workers, num_leaves, costPerUnit, work);

    // Copy to output.
    Tensor* output_finished = nullptr;
    TensorShape finished_shape;
    finished_shape.AddDim(finished_leaves.size());
    OP_REQUIRES_OK(context,
                   context->allocate_output(0, finished_shape,
                                            &output_finished));
    auto out_finished = output_finished->unaligned_flat<int32>();
    std::copy(finished_leaves.begin(), finished_leaves.end(),
              out_finished.data());

    Tensor* output_stale = nullptr;
    TensorShape stale_shape;
    stale_shape.AddDim(stale.size());
    OP_REQUIRES_OK(context,
                   context->allocate_output(1, stale_shape,
                                            &output_stale));
    auto out_stale = output_stale->unaligned_flat<int32>();
    std::copy(stale.begin(), stale.end(), out_stale.data());
  }

 private:
  bool regression_;
  int32 num_split_after_samples_;
  int32 min_split_samples_;
  float dominate_fraction_;
  string dominate_method_;
  int32 random_seed_;
};

REGISTER_KERNEL_BUILDER(Name("FinishedNodes").Device(DEVICE_CPU),
                        FinishedNodes);

}  // namespace tensorflow
<|endoftext|>