<commit_before>// Time:  O(C(n, c)), try out all possible substrings with the minimum c deletion.
// Space: O(c), the depth is at most c, and it costs n at each depth

// DFS solution. (4ms)
class Solution {
public:
    vector<string> removeInvalidParentheses(string s)  {
        int left_removed = 0, right_removed = 0;
        findMinRemove(s, &left_removed, &right_removed);

        vector<string> res;
        vector<int> removed;
        removeInvalidParenthesesHelper(s, 0, left_removed, right_removed, &removed, &res);
        return res;
    }

private:
    void findMinRemove(const string& s, int *left_removed, int *right_removed) {
        // Calculate the minimum left and right parantheses to remove.
        for (const auto& c : s) {
            if (c == '(') {
                ++(*left_removed);
            } else if (c == ')') {
                if (!(*left_removed)) {
                    ++(*right_removed);
                } else {
                    --(*left_removed);
                }
            }
        }
    }

    void removeInvalidParenthesesHelper(const string& s, int start,
                                  int left_removed, int right_removed,
                                  vector<int> *removed, vector<string> *res) {

        if (left_removed == 0 && right_removed == 0) {
            string tmp;
            for (int i = 0, j = 0; i < s.length(); ++i) {
                if (j < removed->size() && i == (*removed)[j]) {
                    ++j;
                } else {
                    tmp.push_back(s[i]);
                }
            }
            if (isValid(tmp)) {
                res->emplace_back(tmp);
            }
            return;
        }

        for (int i = start; i < s.length(); ++i) {
            if (right_removed == 0 && left_removed > 0 && s[i] == '(') {
                if (i == start || s[i] != s[i - 1]) {  // Skip duplicated.
                    removed->emplace_back(i);
                    removeInvalidParenthesesHelper(s, i + 1, left_removed - 1, right_removed,
                                                   removed, res);
                    removed->pop_back();
                }
            } else if (right_removed > 0 && s[i] == ')') {
                if (i == start || s[i] != s[i - 1]) {  // Skip duplicated.
                    removed->emplace_back(i);
                    removeInvalidParenthesesHelper(s, i + 1, left_removed, right_removed - 1,
                                                   removed, res);
                    removed->pop_back();
                }
            }

        }
    }

    // Check whether s is valid or not.
    bool isValid(string s) {
        int sum = 0;
        for (const auto &c : s) {
            if (c == '(') {
                ++sum;
            } else if (c == ')') {
                --sum;
            }
            if (sum < 0) {
                return false;
            }
        }
        return sum == 0;
    }
};

// Time:  O(C(n, c)), try out all possible substrings with the minimum c deletion.
// Space: O(c), the depth is at most c, and it costs n at each depth

// DFS solution. (8ms)
class Solution2 {
public:
    vector<string> removeInvalidParentheses(string s)  {
        int left_removed = 0, right_removed = 0;
        findMinRemove(s, &left_removed, &right_removed);

        vector<string> res;
        unordered_set<int> removed;
        removeInvalidParenthesesHelper(s, 0, left_removed, right_removed, &removed, &res);
        return res;
    }

private:
    void findMinRemove(const string& s, int *left_removed, int *right_removed) {
        // Calculate the minimum left and right parantheses to remove.
        for (const auto& c : s) {
            if (c == '(') {
                ++(*left_removed);
            } else if (c == ')') {
                if (!(*left_removed)) {
                    ++(*right_removed);
                } else {
                    --(*left_removed);
                }
            }
        }
    }

    void removeInvalidParenthesesHelper(const string& s, int start,
                                  int left_removed, int right_removed,
                                  unordered_set<int> *removed, vector<string> *res) {

        if (left_removed == 0 && right_removed == 0) {
            string tmp;
            for (int i = 0; i < s.length(); ++i) {
                if (!removed->count(i)) {
                    tmp.push_back(s[i]);
                }
            }
            if (isValid(tmp)) {
                res->emplace_back(tmp);
            }
            return;
        }

        for (int i = start; i < s.length(); ++i) {
            if (right_removed == 0 && left_removed > 0 && s[i] == '(') {
                if (i == start || s[i] != s[i - 1]) {  // Skip duplicated.
                    removed->emplace(i);
                    removeInvalidParenthesesHelper(s, i + 1, left_removed - 1, right_removed,
                                                   removed, res);
                    removed->erase(i);
                }
            } else if (right_removed > 0 && s[i] == ')') {
                if (i == start || s[i] != s[i - 1]) {  // Skip duplicated.
                    removed->emplace(i);
                    removeInvalidParenthesesHelper(s, i + 1, left_removed, right_removed - 1,
                                                   removed, res);
                    removed->erase(i);
                }
            }

        }
    }

    // Check whether s is valid or not.
    bool isValid(string s) {
        int sum = 0;
        for (const auto &c : s) {
            if (c == '(') {
                ++sum;
            } else if (c == ')') {
                --sum;
            }
            if (sum < 0) {
                return false;
            }
        }
        return sum == 0;
    }
};

// Time:  O(n * C(n, c)), try out all possible substrings with the minimum c deletion.
// Space: O(n * c), the depth is at most c, and it costs n at each depth
// DFS solution. (4ms)
class Solution3 {
public:
    vector<string> removeInvalidParentheses(string s)  {
        int left_removed = 0, right_removed = 0;
        findMinRemove(s, &left_removed, &right_removed);

        vector<string> res;
        removeInvalidParenthesesHelper(s, 0, left_removed, right_removed, &res);
        return res;
    }

    void findMinRemove(const string& s, int *left_removed, int *right_removed) {
        // Calculate the minimum left and right parantheses to remove.
        for (const auto& c : s) {
            if (c == '(') {
                ++(*left_removed);
            } else if (c == ')') {
                if (!(*left_removed)) {
                    ++(*right_removed);
                } else {
                    --(*left_removed);
                }
            }
        }
    }

private:
    void removeInvalidParenthesesHelper(const string& s, int start,
                                  int left_removed, int right_removed, vector<string> *res) {

        if (left_removed == 0 && right_removed == 0) {
            if (isValid(s)) {
                res->emplace_back(s);
            }
            return;
        }

        for (int i = start; i < s.length(); ++i) {
            if (right_removed == 0 && left_removed > 0 && s[i] == '(') {
                if (i == start || s[i] != s[i - 1]) {  // Skip duplicated.
                    string tmp = s;
                    tmp.erase(i, 1);
                    removeInvalidParenthesesHelper(tmp, i, left_removed - 1, right_removed, res);
                }
            } else if (right_removed > 0 && s[i] == ')') {
                if (i == start || s[i] != s[i - 1]) {  // Skip duplicated.
                    string tmp = s;
                    tmp.erase(i, 1);
                    removeInvalidParenthesesHelper(tmp, i, left_removed, right_removed - 1, res);
                }
            }

        }
    }

    // Check whether s is valid or not.
    bool isValid(string s) {
        int sum = 0;
        for (const auto &c : s) {
            if (c == '(') {
                ++sum;
            } else if (c == ')') {
                --sum;
            }
            if (sum < 0) {
                return false;
            }
        }
        return sum == 0;
    }
};
<commit_msg>Update remove-invalid-parentheses.cpp<commit_after>// Time:  O(C(n, c)), try out all possible substrings with the minimum c deletion.
// Space: O(c), the depth is at most c, and it costs n at each depth

// DFS solution with removed array. (4ms)
class Solution {
public:
    vector<string> removeInvalidParentheses(string s)  {
        int left_removed = 0, right_removed = 0;
        findMinRemove(s, &left_removed, &right_removed);

        vector<string> res;
        vector<int> removed;
        removeInvalidParenthesesHelper(s, 0, left_removed, right_removed, &removed, &res);
        return res;
    }

private:
    void findMinRemove(const string& s, int *left_removed, int *right_removed) {
        // Calculate the minimum left and right parantheses to remove.
        for (const auto& c : s) {
            if (c == '(') {
                ++(*left_removed);
            } else if (c == ')') {
                if (!(*left_removed)) {
                    ++(*right_removed);
                } else {
                    --(*left_removed);
                }
            }
        }
    }

    void removeInvalidParenthesesHelper(const string& s, int start,
                                  int left_removed, int right_removed,
                                  vector<int> *removed, vector<string> *res) {

        if (left_removed == 0 && right_removed == 0) {
            string tmp;
            for (int i = 0, j = 0; i < s.length(); ++i) {
                if (j < removed->size() && i == (*removed)[j]) {
                    ++j;
                } else {
                    tmp.push_back(s[i]);
                }
            }
            if (isValid(tmp)) {
                res->emplace_back(tmp);
            }
            return;
        }

        for (int i = start; i < s.length(); ++i) {
            if (right_removed == 0 && left_removed > 0 && s[i] == '(') {
                if (i == start || s[i] != s[i - 1]) {  // Skip duplicated.
                    removed->emplace_back(i);
                    removeInvalidParenthesesHelper(s, i + 1, left_removed - 1, right_removed,
                                                   removed, res);
                    removed->pop_back();
                }
            } else if (right_removed > 0 && s[i] == ')') {
                if (i == start || s[i] != s[i - 1]) {  // Skip duplicated.
                    removed->emplace_back(i);
                    removeInvalidParenthesesHelper(s, i + 1, left_removed, right_removed - 1,
                                                   removed, res);
                    removed->pop_back();
                }
            }

        }
    }

    // Check whether s is valid or not.
    bool isValid(string s) {
        int sum = 0;
        for (const auto &c : s) {
            if (c == '(') {
                ++sum;
            } else if (c == ')') {
                --sum;
            }
            if (sum < 0) {
                return false;
            }
        }
        return sum == 0;
    }
};

// Time:  O(C(n, c)), try out all possible substrings with the minimum c deletion.
// Space: O(c), the depth is at most c, and it costs n at each depth
// DFS solution with removed hash. (8ms)
class Solution2 {
public:
    vector<string> removeInvalidParentheses(string s)  {
        int left_removed = 0, right_removed = 0;
        findMinRemove(s, &left_removed, &right_removed);

        vector<string> res;
        unordered_set<int> removed;
        removeInvalidParenthesesHelper(s, 0, left_removed, right_removed, &removed, &res);
        return res;
    }

private:
    void findMinRemove(const string& s, int *left_removed, int *right_removed) {
        // Calculate the minimum left and right parantheses to remove.
        for (const auto& c : s) {
            if (c == '(') {
                ++(*left_removed);
            } else if (c == ')') {
                if (!(*left_removed)) {
                    ++(*right_removed);
                } else {
                    --(*left_removed);
                }
            }
        }
    }

    void removeInvalidParenthesesHelper(const string& s, int start,
                                  int left_removed, int right_removed,
                                  unordered_set<int> *removed, vector<string> *res) {

        if (left_removed == 0 && right_removed == 0) {
            string tmp;
            for (int i = 0; i < s.length(); ++i) {
                if (!removed->count(i)) {
                    tmp.push_back(s[i]);
                }
            }
            if (isValid(tmp)) {
                res->emplace_back(tmp);
            }
            return;
        }

        for (int i = start; i < s.length(); ++i) {
            if (right_removed == 0 && left_removed > 0 && s[i] == '(') {
                if (i == start || s[i] != s[i - 1]) {  // Skip duplicated.
                    removed->emplace(i);
                    removeInvalidParenthesesHelper(s, i + 1, left_removed - 1, right_removed,
                                                   removed, res);
                    removed->erase(i);
                }
            } else if (right_removed > 0 && s[i] == ')') {
                if (i == start || s[i] != s[i - 1]) {  // Skip duplicated.
                    removed->emplace(i);
                    removeInvalidParenthesesHelper(s, i + 1, left_removed, right_removed - 1,
                                                   removed, res);
                    removed->erase(i);
                }
            }

        }
    }

    // Check whether s is valid or not.
    bool isValid(string s) {
        int sum = 0;
        for (const auto &c : s) {
            if (c == '(') {
                ++sum;
            } else if (c == ')') {
                --sum;
            }
            if (sum < 0) {
                return false;
            }
        }
        return sum == 0;
    }
};

// Time:  O(n * C(n, c)), try out all possible substrings with the minimum c deletion.
// Space: O(n * c), the depth is at most c, and it costs n at each depth
// DFS solution. (4ms)
class Solution3 {
public:
    vector<string> removeInvalidParentheses(string s)  {
        int left_removed = 0, right_removed = 0;
        findMinRemove(s, &left_removed, &right_removed);

        vector<string> res;
        removeInvalidParenthesesHelper(s, 0, left_removed, right_removed, &res);
        return res;
    }

    void findMinRemove(const string& s, int *left_removed, int *right_removed) {
        // Calculate the minimum left and right parantheses to remove.
        for (const auto& c : s) {
            if (c == '(') {
                ++(*left_removed);
            } else if (c == ')') {
                if (!(*left_removed)) {
                    ++(*right_removed);
                } else {
                    --(*left_removed);
                }
            }
        }
    }

private:
    void removeInvalidParenthesesHelper(const string& s, int start,
                                  int left_removed, int right_removed, vector<string> *res) {

        if (left_removed == 0 && right_removed == 0) {
            if (isValid(s)) {
                res->emplace_back(s);
            }
            return;
        }

        for (int i = start; i < s.length(); ++i) {
            if (right_removed == 0 && left_removed > 0 && s[i] == '(') {
                if (i == start || s[i] != s[i - 1]) {  // Skip duplicated.
                    string tmp = s;
                    tmp.erase(i, 1);
                    removeInvalidParenthesesHelper(tmp, i, left_removed - 1, right_removed, res);
                }
            } else if (right_removed > 0 && s[i] == ')') {
                if (i == start || s[i] != s[i - 1]) {  // Skip duplicated.
                    string tmp = s;
                    tmp.erase(i, 1);
                    removeInvalidParenthesesHelper(tmp, i, left_removed, right_removed - 1, res);
                }
            }

        }
    }

    // Check whether s is valid or not.
    bool isValid(string s) {
        int sum = 0;
        for (const auto &c : s) {
            if (c == '(') {
                ++sum;
            } else if (c == ')') {
                --sum;
            }
            if (sum < 0) {
                return false;
            }
        }
        return sum == 0;
    }
};
<|endoftext|>


<commit_before>// Copyright 2008, Google Inc.
// All rights reserved.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are
// met:
//
//    * Redistributions of source code must retain the above copyright
// notice, this list of conditions and the following disclaimer.
//    * Redistributions in binary form must reproduce the above
// copyright notice, this list of conditions and the following disclaimer
// in the documentation and/or other materials provided with the
// distribution.
//    * Neither the name of Google Inc. nor the names of its
// contributors may be used to endorse or promote products derived from
// this software without specific prior written permission.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
// "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
// LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
// A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
// OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
// SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
// LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
// DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
// THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
// (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

#include <string.h>
#include <time.h>
#include <vector>

#include "base/gfx/convolver.h"
#include "testing/gtest/include/gtest/gtest.h"

namespace gfx {

namespace {

// Fills the given filter with impulse functions for the range 0->num_entries.
  void FillImpulseFilter(int num_entries, ConvolusionFilter1D* filter) {
  float one = 1.0f;
  for (int i = 0; i < num_entries; i++)
    filter->AddFilter(i, &one, 1);
}

// Filters the given input with the impulse function, and verifies that it
// does not change.
void TestImpulseConvolusion(const unsigned char* data, int width, int height) {
  int byte_count = width * height * 4;

  ConvolusionFilter1D filter_x;
  FillImpulseFilter(width, &filter_x);

  ConvolusionFilter1D filter_y;
  FillImpulseFilter(height, &filter_y);

  std::vector<unsigned char> output;
  output.resize(byte_count);
  BGRAConvolve2D(data, width * 4, true, filter_x, filter_y, &output[0]);

  // Output should exactly match input.
  EXPECT_EQ(0, memcmp(data, &output[0], byte_count));
}

// Fills the destination filter with a box filter averaging every two pixels
// to produce the output.
void FillBoxFilter(int size, ConvolusionFilter1D* filter) {
  const float box[2] = { 0.5, 0.5 };
  for (int i = 0; i < size; i++)
    filter->AddFilter(i * 2, box, 2);
}

}  // namespace

// Tests that each pixel, when set and run through the impulse filter, does
// not change.
TEST(Convolver, Impulse) {
  // We pick an "odd" size that is not likely to fit on any boundaries so that
  // we can see if all the widths and paddings are handled properly.
  int width = 15;
  int height = 31;
  int byte_count = width * height * 4;
  std::vector<unsigned char> input;
  input.resize(byte_count);

  unsigned char* input_ptr = &input[0];
  for (int y = 0; y < height; y++) {
    for (int x = 0; x < width; x++) {
      for (int channel = 0; channel < 3; channel++) {
        memset(input_ptr, 0, byte_count);
        input_ptr[(y * width + x) * 4 + channel] = 0xff;
        // Always set the alpha channel or it will attempt to "fix" it for us.
        input_ptr[(y * width + x) * 4 + 3] = 0xff;
        TestImpulseConvolusion(input_ptr, width, height);
      }
    }
  }
}

// Tests that using a box filter to halve an image results in every square of 4
// pixels in the original get averaged to a pixel in the output.
TEST(Convolver, Halve) {
  static const int kSize = 16;

  int src_width = kSize;
  int src_height = kSize;
  int src_row_stride = src_width * 4;
  int src_byte_count = src_row_stride * src_height;
  std::vector<unsigned char> input;
  input.resize(src_byte_count);

  int dest_width = src_width / 2;
  int dest_height = src_height / 2;
  int dest_byte_count = dest_width * dest_height * 4;
  std::vector<unsigned char> output;
  output.resize(dest_byte_count);

  // First fill the array with a bunch of random data.
  srand(static_cast<unsigned>(time(NULL)));
  for (int i = 0; i < src_byte_count; i++)
    input[i] = rand() * 255 / RAND_MAX;

  // Compute the filters.
  ConvolusionFilter1D filter_x, filter_y;
  FillBoxFilter(dest_width, &filter_x);
  FillBoxFilter(dest_height, &filter_y);

  // Do the convolusion.
  BGRAConvolve2D(&input[0], src_width, true, filter_x, filter_y, &output[0]);

  // Compute the expected results and check, allowing for a small difference
  // to account for rounding errors.
  for (int y = 0; y < dest_height; y++) {
    for (int x = 0; x < dest_width; x++) {
      for (int channel = 0; channel < 4; channel++) {
        int src_offset = (y * 2 * src_row_stride + x * 2 * 4) + channel;
        int value = input[src_offset] +  // Top left source pixel.
                    input[src_offset + 4] +  // Top right source pixel.
                    input[src_offset + src_row_stride] +  // Lower left.
                    input[src_offset + src_row_stride + 4];  // Lower right.
        value /= 4;  // Average.
        int difference = value - output[(y * dest_width + x) * 4 + channel];
        EXPECT_TRUE(difference >= -1 || difference <= 1);
      }
    }
  }
}

}  // namespace gfx<commit_msg>Add a EOL to the end of convolver_unittest.cc.  This eliminates a GCC warning.<commit_after>// Copyright 2008, Google Inc.
// All rights reserved.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are
// met:
//
//    * Redistributions of source code must retain the above copyright
// notice, this list of conditions and the following disclaimer.
//    * Redistributions in binary form must reproduce the above
// copyright notice, this list of conditions and the following disclaimer
// in the documentation and/or other materials provided with the
// distribution.
//    * Neither the name of Google Inc. nor the names of its
// contributors may be used to endorse or promote products derived from
// this software without specific prior written permission.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
// "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
// LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
// A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
// OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
// SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
// LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
// DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
// THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
// (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

#include <string.h>
#include <time.h>
#include <vector>

#include "base/gfx/convolver.h"
#include "testing/gtest/include/gtest/gtest.h"

namespace gfx {

namespace {

// Fills the given filter with impulse functions for the range 0->num_entries.
  void FillImpulseFilter(int num_entries, ConvolusionFilter1D* filter) {
  float one = 1.0f;
  for (int i = 0; i < num_entries; i++)
    filter->AddFilter(i, &one, 1);
}

// Filters the given input with the impulse function, and verifies that it
// does not change.
void TestImpulseConvolusion(const unsigned char* data, int width, int height) {
  int byte_count = width * height * 4;

  ConvolusionFilter1D filter_x;
  FillImpulseFilter(width, &filter_x);

  ConvolusionFilter1D filter_y;
  FillImpulseFilter(height, &filter_y);

  std::vector<unsigned char> output;
  output.resize(byte_count);
  BGRAConvolve2D(data, width * 4, true, filter_x, filter_y, &output[0]);

  // Output should exactly match input.
  EXPECT_EQ(0, memcmp(data, &output[0], byte_count));
}

// Fills the destination filter with a box filter averaging every two pixels
// to produce the output.
void FillBoxFilter(int size, ConvolusionFilter1D* filter) {
  const float box[2] = { 0.5, 0.5 };
  for (int i = 0; i < size; i++)
    filter->AddFilter(i * 2, box, 2);
}

}  // namespace

// Tests that each pixel, when set and run through the impulse filter, does
// not change.
TEST(Convolver, Impulse) {
  // We pick an "odd" size that is not likely to fit on any boundaries so that
  // we can see if all the widths and paddings are handled properly.
  int width = 15;
  int height = 31;
  int byte_count = width * height * 4;
  std::vector<unsigned char> input;
  input.resize(byte_count);

  unsigned char* input_ptr = &input[0];
  for (int y = 0; y < height; y++) {
    for (int x = 0; x < width; x++) {
      for (int channel = 0; channel < 3; channel++) {
        memset(input_ptr, 0, byte_count);
        input_ptr[(y * width + x) * 4 + channel] = 0xff;
        // Always set the alpha channel or it will attempt to "fix" it for us.
        input_ptr[(y * width + x) * 4 + 3] = 0xff;
        TestImpulseConvolusion(input_ptr, width, height);
      }
    }
  }
}

// Tests that using a box filter to halve an image results in every square of 4
// pixels in the original get averaged to a pixel in the output.
TEST(Convolver, Halve) {
  static const int kSize = 16;

  int src_width = kSize;
  int src_height = kSize;
  int src_row_stride = src_width * 4;
  int src_byte_count = src_row_stride * src_height;
  std::vector<unsigned char> input;
  input.resize(src_byte_count);

  int dest_width = src_width / 2;
  int dest_height = src_height / 2;
  int dest_byte_count = dest_width * dest_height * 4;
  std::vector<unsigned char> output;
  output.resize(dest_byte_count);

  // First fill the array with a bunch of random data.
  srand(static_cast<unsigned>(time(NULL)));
  for (int i = 0; i < src_byte_count; i++)
    input[i] = rand() * 255 / RAND_MAX;

  // Compute the filters.
  ConvolusionFilter1D filter_x, filter_y;
  FillBoxFilter(dest_width, &filter_x);
  FillBoxFilter(dest_height, &filter_y);

  // Do the convolusion.
  BGRAConvolve2D(&input[0], src_width, true, filter_x, filter_y, &output[0]);

  // Compute the expected results and check, allowing for a small difference
  // to account for rounding errors.
  for (int y = 0; y < dest_height; y++) {
    for (int x = 0; x < dest_width; x++) {
      for (int channel = 0; channel < 4; channel++) {
        int src_offset = (y * 2 * src_row_stride + x * 2 * 4) + channel;
        int value = input[src_offset] +  // Top left source pixel.
                    input[src_offset + 4] +  // Top right source pixel.
                    input[src_offset + src_row_stride] +  // Lower left.
                    input[src_offset + src_row_stride + 4];  // Lower right.
        value /= 4;  // Average.
        int difference = value - output[(y * dest_width + x) * 4 + channel];
        EXPECT_TRUE(difference >= -1 || difference <= 1);
      }
    }
  }
}

}  // namespace gfx
<|endoftext|>


<commit_before>#include "rapid_perception/pose_estimation.h"

#include <algorithm>
#include <cmath>
#include <iostream>
#include <limits.h>
#include <string>
#include <utility>
#include <vector>

#include "Eigen/Core"
#include "Eigen/Geometry"
#include "geometry_msgs/Vector3.h"
#include "pcl/PointIndices.h"
#include "pcl/common/common.h"
#include "pcl/filters/extract_indices.h"
#include "pcl/point_cloud.h"
#include "pcl/point_types.h"
#include "pcl/registration/icp.h"
#include "pcl/search/kdtree.h"
#include "pcl_conversions/pcl_conversions.h"
#include "sensor_msgs/PointCloud2.h"

#include "rapid_msgs/Roi3D.h"
#include "rapid_perception/icp_fitness_functions.h"
#include "rapid_perception/template_matching_heat_mapper.h"
#include "rapid_utils/stochastic_universal_sampling.h"
#include "rapid_viz/publish.h"

typedef pcl::PointXYZRGB PointC;
typedef pcl::PointXYZ PointP;
typedef pcl::PointCloud<PointC> PointCloudC;
typedef pcl::PointCloud<PointP> PointCloudP;
typedef pcl::FPFHSignature33 FPFH;
typedef pcl::PointCloud<FPFH> PointCloudF;
typedef pcl::search::KdTree<FPFH> FeatureTree;
typedef pcl::search::KdTree<PointC> PointCTree;
typedef pcl::search::KdTree<PointP> PointPTree;
using std::pair;
using std::string;
using std::vector;

namespace rapid {
namespace perception {
PoseEstimationMatch::PoseEstimationMatch(PointCloudC::Ptr cloud, double fitness)
    : cloud_(new PointCloudC), fitness_(fitness) {
  *cloud_ = *cloud;
  Eigen::Vector4f min_pt;
  Eigen::Vector4f max_pt;
  pcl::getMinMax3D(*cloud_, min_pt, max_pt);
  center_.x = min_pt.x() + (max_pt.x() - min_pt.x()) / 2;
  center_.y = min_pt.y() + (max_pt.y() - min_pt.y()) / 2;
  center_.z = min_pt.z() + (max_pt.z() - min_pt.z()) / 2;
}

pcl::PointCloud<pcl::PointXYZRGB>::Ptr PoseEstimationMatch::cloud() {
  return cloud_;
}

pcl::PointXYZ PoseEstimationMatch::center() const { return center_; }

double PoseEstimationMatch::fitness() const { return fitness_; }

void PoseEstimationMatch::set_fitness(double fitness) { fitness_ = fitness; }

PoseEstimator::PoseEstimator()
    : scene_(new PointCloudC()),
      object_(new PointCloudC()),
      object_center_(),
      num_candidates_(100),
      fitness_threshold_(0.00002),
      nms_radius_(0.03),
      debug_(false),
      candidates_pub_(),
      alignment_pub_(),
      output_pub_() {}

void PoseEstimator::set_heat_mapper(PoseEstimationHeatMapper* heat_mapper) {
  heat_mapper_ = heat_mapper;
}

void PoseEstimator::set_scene(PointCloudC::Ptr scene) {
  scene_ = scene;
  heat_mapper_->set_scene(scene);
}

void PoseEstimator::set_object(
    const pcl::PointCloud<pcl::PointXYZRGB>::Ptr& object,
    const rapid_msgs::Roi3D& roi) {
  object_ = object;
  object_roi_ = roi;

  // Estimate object radius, just average of x/y/z dimensions.
  Eigen::Vector4f min_pt;
  Eigen::Vector4f max_pt;
  pcl::getMinMax3D(*object_, min_pt, max_pt);
  object_center_.x() = (max_pt.x() + min_pt.x()) / 2;
  object_center_.y() = (max_pt.y() + min_pt.y()) / 2;
  object_center_.z() = (max_pt.z() + min_pt.z()) / 2;

  heat_mapper_->set_object(object_);
  if (heat_mapper_->name() == "template_matching") {
    static_cast<TemplateMatchingHeatMapper*>(heat_mapper_)->set_object_roi(roi);
  }
}

PoseEstimationHeatMapper* PoseEstimator::heat_mapper() { return heat_mapper_; }

void PoseEstimator::set_debug(bool val) {
  debug_ = val;
  heat_mapper_->set_debug(val);
}
void PoseEstimator::set_num_candidates(int val) { num_candidates_ = val; }
void PoseEstimator::set_fitness_threshold(double val) {
  fitness_threshold_ = val;
}
void PoseEstimator::set_nms_radius(double val) { nms_radius_ = val; }

void PoseEstimator::set_candidates_publisher(const ros::Publisher& pub) {
  candidates_pub_ = pub;
}
void PoseEstimator::set_alignment_publisher(const ros::Publisher& pub) {
  alignment_pub_ = pub;
}
void PoseEstimator::set_output_publisher(const ros::Publisher& pub) {
  output_pub_ = pub;
}

bool PoseEstimator::Find() {
  // Compute heatmap of features
  pcl::PointIndicesPtr heatmap_indices(new pcl::PointIndices);
  Eigen::VectorXd importances;
  heat_mapper_->Compute(heatmap_indices, &importances);

  // Sample candidate points from the heatmap
  pcl::PointIndices::Ptr candidate_indices(new pcl::PointIndices);
  ComputeTopCandidates(importances, heatmap_indices, candidate_indices);

  // For each candidate point, run ICP
  vector<PoseEstimationMatch> output_objects;
  RunIcpCandidates(candidate_indices, &output_objects);
  if (output_objects.size() == 0) {
    ROS_WARN("No matches found.");
    return false;
  }

  // Non-max suppression
  vector<bool> keep;
  NonMaxSuppression(output_objects, &keep);

  PointCloudC::Ptr output_cloud(new PointCloudC);
  int num_instances = 0;
  for (size_t i = 0; i < output_objects.size(); ++i) {
    if (keep[i]) {
      double r = static_cast<double>(rand()) / RAND_MAX;
      double g = static_cast<double>(rand()) / RAND_MAX;
      double b = static_cast<double>(rand()) / RAND_MAX;
      Colorize(output_objects[i].cloud(), r, g, b);
      *output_cloud += *output_objects[i].cloud();
      ++num_instances;
    }
  }
  ROS_INFO("Found %d instances of the model", num_instances);
  output_cloud->header.frame_id = scene_->header.frame_id;
  viz::PublishCloud(output_pub_, *output_cloud);
  return true;
}

void PoseEstimator::ComputeCandidates(Eigen::VectorXd& importances,
                                      pcl::PointIndicesPtr heatmap_indices,
                                      pcl::PointIndicesPtr candidate_indices) {
  // Sample points based on importance
  double sum = importances.sum();
  importances /= sum;

  vector<double> cdf;
  double cdf_val = 0;
  for (size_t indices_i = 0; indices_i < heatmap_indices->indices.size();
       ++indices_i) {
    double prob = importances(indices_i);
    cdf_val += prob;
    cdf.push_back(cdf_val);
  }

  // Indices of indices! We are sampling the indices vector by importance. The
  // sampled indices themselves are stored in sampled_important_indices, so the
  // actual point index for i in sampled_important_indices is
  // cloud[indices[sampled_important_indices[i]]].
  vector<int> sampled_important_indices;
  utils::StochasticUniversalSampling(cdf, num_candidates_,
                                     &sampled_important_indices);

  for (size_t i = 0; i < sampled_important_indices.size(); ++i) {
    int index_into_indices = sampled_important_indices[i];
    int point_index = heatmap_indices->indices[index_into_indices];
    candidate_indices->indices.push_back(point_index);
  }

  // Visualize the candidate points
  PointCloudC::Ptr viz_cloud(new PointCloudC());
  viz_cloud.reset(new PointCloudC());
  pcl::ExtractIndices<PointC> extract;
  extract.setInputCloud(scene_);
  extract.setIndices(candidate_indices);
  extract.filter(*viz_cloud);
  viz::PublishCloud(candidates_pub_, *viz_cloud);
}

// This is used to sort heatmap scores, represented as (index, score).
// Higher scores are better.
bool CompareScores(const pair<int, double>& a, const pair<int, double>& b) {
  return a.second > b.second;
}

void PoseEstimator::ComputeTopCandidates(
    Eigen::VectorXd& importances, pcl::PointIndicesPtr heatmap_indices,
    pcl::PointIndicesPtr candidate_indices) {
  std::vector<std::pair<int, double> > scores;
  for (size_t i = 0; i < heatmap_indices->indices.size(); ++i) {
    int index = heatmap_indices->indices[i];
    double score = importances[i];
    scores.push_back(std::make_pair(index, score));
  }

  size_t num = std::min(scores.size(), static_cast<size_t>(num_candidates_));
  std::partial_sort(scores.begin(), scores.begin() + num, scores.end(),
                    &CompareScores);

  candidate_indices->indices.clear();
  for (size_t i = 0; i < num; ++i) {
    candidate_indices->indices.push_back(scores[i].first);
  }

  // Visualize the candidate points
  PointCloudC::Ptr viz_cloud(new PointCloudC());
  viz_cloud.reset(new PointCloudC());
  pcl::ExtractIndices<PointC> extract;
  extract.setInputCloud(scene_);
  extract.setIndices(candidate_indices);
  extract.filter(*viz_cloud);
  viz::PublishCloud(candidates_pub_, *viz_cloud);
}

void PoseEstimator::RunIcpCandidates(
    pcl::PointIndices::Ptr candidate_indices,
    vector<PoseEstimationMatch>* output_objects) {
  // Copy the object and its center.
  PointCloudC::Ptr working_object(new PointCloudC);
  PointCloudC::Ptr aligned_object(new PointCloudC);
  pcl::IterativeClosestPoint<PointC, PointC> icp;
  icp.setInputTarget(scene_);
  output_objects->clear();
  for (size_t ci = 0; ci < candidate_indices->indices.size(); ++ci) {
    int index = candidate_indices->indices[ci];
    const PointC& center = scene_->points[index];
    // Transform the object to be centered on the center point
    // So far we just use a translation offset
    // Also assume that the object and scene are in the same frame
    Eigen::Vector3d translation;
    translation.x() = center.x - object_center_.x();
    translation.y() = center.y - object_center_.y();
    translation.z() = center.z - object_center_.z();
    pcl::transformPointCloud(*object_, *working_object, translation,
                             Eigen::Quaterniond::Identity());

    // Run ICP
    icp.setInputSource(working_object);
    icp.align(*aligned_object);
    viz::PublishCloud(alignment_pub_, *aligned_object);
    if (icp.hasConverged()) {
      Eigen::Matrix4f final_transform = icp.getFinalTransformation();
      rapid_msgs::Roi3D roi = object_roi_;
      roi.transform.translation.x += translation.x();
      roi.transform.translation.y += translation.y();
      roi.transform.translation.z += translation.z();
      Eigen::Matrix3f rot_matrix(final_transform.topLeftCorner(3, 3));
      Eigen::Quaternionf q(rot_matrix);
      roi.transform.rotation.w = q.w();
      roi.transform.rotation.x = q.x();
      roi.transform.rotation.y = q.y();
      roi.transform.rotation.z = q.z();

      double fitness = ComputeIcpFitness(scene_, aligned_object, roi);

      string threshold_marker = "";
      if (fitness < fitness_threshold_) {
        output_objects->push_back(PoseEstimationMatch(aligned_object, fitness));
        threshold_marker = " -- below threshold";
      }
      ROS_INFO("ICP converged with score: %f%s", fitness,
               threshold_marker.c_str());
    }
    if (debug_) {
      std::cout << "Press enter to continue ";
      string user_input;
      std::getline(std::cin, user_input);
    }
  }
}

void PoseEstimator::NonMaxSuppression(
    vector<PoseEstimationMatch>& output_objects, vector<bool>* keep) {
  // Non-max supression of output objects
  // Index centers into a tree
  PointCloudP::Ptr output_centers(new PointCloudP);
  for (size_t i = 0; i < output_objects.size(); ++i) {
    output_centers->push_back(output_objects[i].center());
  }
  PointPTree output_centers_tree;
  output_centers_tree.setInputCloud(output_centers);

  // Suppress non max within a radius
  for (int i = 0; i < static_cast<int>(output_objects.size()); ++i) {
    PoseEstimationMatch& match = output_objects[i];
    vector<int> neighbor_indices;
    vector<float> neighbor_distances;
    output_centers_tree.radiusSearch(match.center(), nms_radius_,
                                     neighbor_indices, neighbor_distances);

    bool keep_match = true;
    for (size_t k = 0; k < neighbor_indices.size(); ++k) {
      int index = neighbor_indices[k];
      if (index == i) {
        continue;
      }
      PoseEstimationMatch& neighbor = output_objects[index];
      if (neighbor.fitness() < match.fitness()) {
        if (debug_) {
          ROS_INFO("Suppressing match %d, %f >= neighbor %d, %f", i,
                   match.fitness(), index, neighbor.fitness());
        }
        keep_match = false;
        break;
      } else if (neighbor.fitness() == match.fitness()) {
        // Keep this one, bump neighbor's fitness score up so that it won't come
        // up again.
        neighbor.set_fitness(neighbor.fitness() * 1.01);
      }
    }
    keep->push_back(keep_match);
  }
}

void Colorize(pcl::PointCloud<pcl::PointXYZRGB>::Ptr cloud, double r, double g,
              double b) {
  for (size_t i = 0; i < cloud->size(); ++i) {
    cloud->points[i].r = static_cast<int>(round(r * 255));
    cloud->points[i].g = static_cast<int>(round(g * 255));
    cloud->points[i].b = static_cast<int>(round(b * 255));
  }
}
}  // namespace perception
}  // namespace rapid
<commit_msg>Report best fitness score seen.<commit_after>#include "rapid_perception/pose_estimation.h"

#include <algorithm>
#include <cmath>
#include <iostream>
#include <limits.h>
#include <string>
#include <utility>
#include <vector>

#include "Eigen/Core"
#include "Eigen/Geometry"
#include "geometry_msgs/Vector3.h"
#include "pcl/PointIndices.h"
#include "pcl/common/common.h"
#include "pcl/filters/extract_indices.h"
#include "pcl/point_cloud.h"
#include "pcl/point_types.h"
#include "pcl/registration/icp.h"
#include "pcl/search/kdtree.h"
#include "pcl_conversions/pcl_conversions.h"
#include "sensor_msgs/PointCloud2.h"

#include "rapid_msgs/Roi3D.h"
#include "rapid_perception/icp_fitness_functions.h"
#include "rapid_perception/template_matching_heat_mapper.h"
#include "rapid_utils/stochastic_universal_sampling.h"
#include "rapid_viz/publish.h"

typedef pcl::PointXYZRGB PointC;
typedef pcl::PointXYZ PointP;
typedef pcl::PointCloud<PointC> PointCloudC;
typedef pcl::PointCloud<PointP> PointCloudP;
typedef pcl::FPFHSignature33 FPFH;
typedef pcl::PointCloud<FPFH> PointCloudF;
typedef pcl::search::KdTree<FPFH> FeatureTree;
typedef pcl::search::KdTree<PointC> PointCTree;
typedef pcl::search::KdTree<PointP> PointPTree;
using std::pair;
using std::string;
using std::vector;

namespace rapid {
namespace perception {
PoseEstimationMatch::PoseEstimationMatch(PointCloudC::Ptr cloud, double fitness)
    : cloud_(new PointCloudC), fitness_(fitness) {
  *cloud_ = *cloud;
  Eigen::Vector4f min_pt;
  Eigen::Vector4f max_pt;
  pcl::getMinMax3D(*cloud_, min_pt, max_pt);
  center_.x = min_pt.x() + (max_pt.x() - min_pt.x()) / 2;
  center_.y = min_pt.y() + (max_pt.y() - min_pt.y()) / 2;
  center_.z = min_pt.z() + (max_pt.z() - min_pt.z()) / 2;
}

pcl::PointCloud<pcl::PointXYZRGB>::Ptr PoseEstimationMatch::cloud() {
  return cloud_;
}

pcl::PointXYZ PoseEstimationMatch::center() const { return center_; }

double PoseEstimationMatch::fitness() const { return fitness_; }

void PoseEstimationMatch::set_fitness(double fitness) { fitness_ = fitness; }

PoseEstimator::PoseEstimator()
    : scene_(new PointCloudC()),
      object_(new PointCloudC()),
      object_center_(),
      num_candidates_(100),
      fitness_threshold_(0.00002),
      nms_radius_(0.03),
      debug_(false),
      candidates_pub_(),
      alignment_pub_(),
      output_pub_() {}

void PoseEstimator::set_heat_mapper(PoseEstimationHeatMapper* heat_mapper) {
  heat_mapper_ = heat_mapper;
}

void PoseEstimator::set_scene(PointCloudC::Ptr scene) {
  scene_ = scene;
  heat_mapper_->set_scene(scene);
}

void PoseEstimator::set_object(
    const pcl::PointCloud<pcl::PointXYZRGB>::Ptr& object,
    const rapid_msgs::Roi3D& roi) {
  object_ = object;
  object_roi_ = roi;

  // Estimate object radius, just average of x/y/z dimensions.
  Eigen::Vector4f min_pt;
  Eigen::Vector4f max_pt;
  pcl::getMinMax3D(*object_, min_pt, max_pt);
  object_center_.x() = (max_pt.x() + min_pt.x()) / 2;
  object_center_.y() = (max_pt.y() + min_pt.y()) / 2;
  object_center_.z() = (max_pt.z() + min_pt.z()) / 2;

  heat_mapper_->set_object(object_);
  if (heat_mapper_->name() == "template_matching") {
    static_cast<TemplateMatchingHeatMapper*>(heat_mapper_)->set_object_roi(roi);
  }
}

PoseEstimationHeatMapper* PoseEstimator::heat_mapper() { return heat_mapper_; }

void PoseEstimator::set_debug(bool val) {
  debug_ = val;
  heat_mapper_->set_debug(val);
}
void PoseEstimator::set_num_candidates(int val) { num_candidates_ = val; }
void PoseEstimator::set_fitness_threshold(double val) {
  fitness_threshold_ = val;
}
void PoseEstimator::set_nms_radius(double val) { nms_radius_ = val; }

void PoseEstimator::set_candidates_publisher(const ros::Publisher& pub) {
  candidates_pub_ = pub;
}
void PoseEstimator::set_alignment_publisher(const ros::Publisher& pub) {
  alignment_pub_ = pub;
}
void PoseEstimator::set_output_publisher(const ros::Publisher& pub) {
  output_pub_ = pub;
}

bool PoseEstimator::Find() {
  // Compute heatmap of features
  pcl::PointIndicesPtr heatmap_indices(new pcl::PointIndices);
  Eigen::VectorXd importances;
  heat_mapper_->Compute(heatmap_indices, &importances);

  // Sample candidate points from the heatmap
  pcl::PointIndices::Ptr candidate_indices(new pcl::PointIndices);
  ComputeTopCandidates(importances, heatmap_indices, candidate_indices);

  // For each candidate point, run ICP
  vector<PoseEstimationMatch> output_objects;
  RunIcpCandidates(candidate_indices, &output_objects);
  if (output_objects.size() == 0) {
    ROS_WARN("No matches found.");
    return false;
  }

  // Non-max suppression
  vector<bool> keep;
  NonMaxSuppression(output_objects, &keep);

  PointCloudC::Ptr output_cloud(new PointCloudC);
  int num_instances = 0;
  for (size_t i = 0; i < output_objects.size(); ++i) {
    if (keep[i]) {
      double r = static_cast<double>(rand()) / RAND_MAX;
      double g = static_cast<double>(rand()) / RAND_MAX;
      double b = static_cast<double>(rand()) / RAND_MAX;
      Colorize(output_objects[i].cloud(), r, g, b);
      *output_cloud += *output_objects[i].cloud();
      ++num_instances;
    }
  }
  ROS_INFO("Found %d instances of the model", num_instances);
  output_cloud->header.frame_id = scene_->header.frame_id;
  viz::PublishCloud(output_pub_, *output_cloud);
  return true;
}

void PoseEstimator::ComputeCandidates(Eigen::VectorXd& importances,
                                      pcl::PointIndicesPtr heatmap_indices,
                                      pcl::PointIndicesPtr candidate_indices) {
  // Sample points based on importance
  double sum = importances.sum();
  importances /= sum;

  vector<double> cdf;
  double cdf_val = 0;
  for (size_t indices_i = 0; indices_i < heatmap_indices->indices.size();
       ++indices_i) {
    double prob = importances(indices_i);
    cdf_val += prob;
    cdf.push_back(cdf_val);
  }

  // Indices of indices! We are sampling the indices vector by importance. The
  // sampled indices themselves are stored in sampled_important_indices, so the
  // actual point index for i in sampled_important_indices is
  // cloud[indices[sampled_important_indices[i]]].
  vector<int> sampled_important_indices;
  utils::StochasticUniversalSampling(cdf, num_candidates_,
                                     &sampled_important_indices);

  for (size_t i = 0; i < sampled_important_indices.size(); ++i) {
    int index_into_indices = sampled_important_indices[i];
    int point_index = heatmap_indices->indices[index_into_indices];
    candidate_indices->indices.push_back(point_index);
  }

  // Visualize the candidate points
  PointCloudC::Ptr viz_cloud(new PointCloudC());
  viz_cloud.reset(new PointCloudC());
  pcl::ExtractIndices<PointC> extract;
  extract.setInputCloud(scene_);
  extract.setIndices(candidate_indices);
  extract.filter(*viz_cloud);
  viz::PublishCloud(candidates_pub_, *viz_cloud);
}

// This is used to sort heatmap scores, represented as (index, score).
// Higher scores are better.
bool CompareScores(const pair<int, double>& a, const pair<int, double>& b) {
  return a.second > b.second;
}

void PoseEstimator::ComputeTopCandidates(
    Eigen::VectorXd& importances, pcl::PointIndicesPtr heatmap_indices,
    pcl::PointIndicesPtr candidate_indices) {
  std::vector<std::pair<int, double> > scores;
  for (size_t i = 0; i < heatmap_indices->indices.size(); ++i) {
    int index = heatmap_indices->indices[i];
    double score = importances[i];
    scores.push_back(std::make_pair(index, score));
  }

  size_t num = std::min(scores.size(), static_cast<size_t>(num_candidates_));
  std::partial_sort(scores.begin(), scores.begin() + num, scores.end(),
                    &CompareScores);

  candidate_indices->indices.clear();
  for (size_t i = 0; i < num; ++i) {
    candidate_indices->indices.push_back(scores[i].first);
  }

  // Visualize the candidate points
  PointCloudC::Ptr viz_cloud(new PointCloudC());
  viz_cloud.reset(new PointCloudC());
  pcl::ExtractIndices<PointC> extract;
  extract.setInputCloud(scene_);
  extract.setIndices(candidate_indices);
  extract.filter(*viz_cloud);
  viz::PublishCloud(candidates_pub_, *viz_cloud);
}

void PoseEstimator::RunIcpCandidates(
    pcl::PointIndices::Ptr candidate_indices,
    vector<PoseEstimationMatch>* output_objects) {
  // Copy the object and its center.
  PointCloudC::Ptr working_object(new PointCloudC);
  PointCloudC::Ptr aligned_object(new PointCloudC);
  pcl::IterativeClosestPoint<PointC, PointC> icp;
  icp.setInputTarget(scene_);
  output_objects->clear();
  double best_fitness = std::numeric_limits<double>::max();
  for (size_t ci = 0; ci < candidate_indices->indices.size(); ++ci) {
    int index = candidate_indices->indices[ci];
    const PointC& center = scene_->points[index];
    // Transform the object to be centered on the center point
    // So far we just use a translation offset
    // Also assume that the object and scene are in the same frame
    Eigen::Vector3d translation;
    translation.x() = center.x - object_center_.x();
    translation.y() = center.y - object_center_.y();
    translation.z() = center.z - object_center_.z();
    pcl::transformPointCloud(*object_, *working_object, translation,
                             Eigen::Quaterniond::Identity());

    // Run ICP
    icp.setInputSource(working_object);
    icp.align(*aligned_object);
    viz::PublishCloud(alignment_pub_, *aligned_object);
    if (icp.hasConverged()) {
      Eigen::Matrix4f final_transform = icp.getFinalTransformation();
      rapid_msgs::Roi3D roi = object_roi_;
      roi.transform.translation.x += translation.x();
      roi.transform.translation.y += translation.y();
      roi.transform.translation.z += translation.z();
      Eigen::Matrix3f rot_matrix(final_transform.topLeftCorner(3, 3));
      Eigen::Quaternionf q(rot_matrix);
      roi.transform.rotation.w = q.w();
      roi.transform.rotation.x = q.x();
      roi.transform.rotation.y = q.y();
      roi.transform.rotation.z = q.z();

      double fitness = ComputeIcpFitness(scene_, aligned_object, roi);
      if (fitness < best_fitness) {
        best_fitness = fitness;
      }

      string threshold_marker = "";
      if (fitness < fitness_threshold_) {
        output_objects->push_back(PoseEstimationMatch(aligned_object, fitness));
        threshold_marker = " -- below threshold";
      }
      ROS_INFO("ICP converged with score: %f%s", fitness,
               threshold_marker.c_str());
    }
    if (debug_) {
      std::cout << "Press enter to continue ";
      string user_input;
      std::getline(std::cin, user_input);
    }
  }
  ROS_INFO("Best score in candidates: %f", best_fitness);
}

void PoseEstimator::NonMaxSuppression(
    vector<PoseEstimationMatch>& output_objects, vector<bool>* keep) {
  // Non-max supression of output objects
  // Index centers into a tree
  PointCloudP::Ptr output_centers(new PointCloudP);
  for (size_t i = 0; i < output_objects.size(); ++i) {
    output_centers->push_back(output_objects[i].center());
  }
  PointPTree output_centers_tree;
  output_centers_tree.setInputCloud(output_centers);

  // Suppress non max within a radius
  for (int i = 0; i < static_cast<int>(output_objects.size()); ++i) {
    PoseEstimationMatch& match = output_objects[i];
    vector<int> neighbor_indices;
    vector<float> neighbor_distances;
    output_centers_tree.radiusSearch(match.center(), nms_radius_,
                                     neighbor_indices, neighbor_distances);

    bool keep_match = true;
    for (size_t k = 0; k < neighbor_indices.size(); ++k) {
      int index = neighbor_indices[k];
      if (index == i) {
        continue;
      }
      PoseEstimationMatch& neighbor = output_objects[index];
      if (neighbor.fitness() < match.fitness()) {
        if (debug_) {
          ROS_INFO("Suppressing match %d, %f >= neighbor %d, %f", i,
                   match.fitness(), index, neighbor.fitness());
        }
        keep_match = false;
        break;
      } else if (neighbor.fitness() == match.fitness()) {
        // Keep this one, bump neighbor's fitness score up so that it won't come
        // up again.
        neighbor.set_fitness(neighbor.fitness() * 1.01);
      }
    }
    keep->push_back(keep_match);
  }
}

void Colorize(pcl::PointCloud<pcl::PointXYZRGB>::Ptr cloud, double r, double g,
              double b) {
  for (size_t i = 0; i < cloud->size(); ++i) {
    cloud->points[i].r = static_cast<int>(round(r * 255));
    cloud->points[i].g = static_cast<int>(round(g * 255));
    cloud->points[i].b = static_cast<int>(round(b * 255));
  }
}
}  // namespace perception
}  // namespace rapid
<|endoftext|>


<commit_before>/*
 * Copyright 2016, Masayoshi Mizutani, mizutani@sfc.wide.ad.jp
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are
 * met:
 *
 * * Redistributions of source code must retain the above copyright
 * notice, this list of conditions and the following disclaimer.
 * * Redistributions in binary form must reproduce the above
 * copyright notice, this list of conditions and the following disclaimer
 * in the documentation and/or other materials provided with the
 * distribution.
 * * Neither the name of Google Inc. nor the names of its
 * contributors may be used to endorse or promote products derived from
 * this software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
 * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
 * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
 * A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
 * OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
 * SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
 * LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
 * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
 * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 */

#include <vector>
#include <string>

#include "./gtest.h"
#include "../argparse.hpp"

TEST(Parser, nargs1) {
  argparse::Parser *psr = new argparse::Parser("test");
  psr->add_argument("-a").nargs(1);
  argparse::Argv args1 = {"./test", "-a", "v"};
  argparse::Argv args2 = {"./test", "-a", "v", "w"};

  argparse::Values val = psr->parse_args(args1);
  EXPECT_EQ(val.get("a"), "v");
  EXPECT_THROW(val.get("a", 1), argparse::exception::IndexError);

  EXPECT_THROW(psr->parse_args(args2), argparse::exception::ParseError);
}

TEST(Parser, nargs2) {
  argparse::Parser *psr = new argparse::Parser("test");
  psr->add_argument("-a").nargs(2);
  argparse::Argv args1 = {"./test", "-a", "v"};
  argparse::Argv args2 = {"./test", "-a", "v", "w"};
  argparse::Argv args3 = {"./test", "-a", "v", "w", "x"};

  argparse::Values val = psr->parse_args(args2);
  EXPECT_EQ(val.get("a"), "v");
  EXPECT_EQ(val.get("a", 0), "v");
  EXPECT_EQ(val.get("a", 1), "w");
  EXPECT_THROW(val.get("a", 2), argparse::exception::IndexError);

  EXPECT_THROW(psr->parse_args(args1), argparse::exception::ParseError);
  EXPECT_THROW(psr->parse_args(args3), argparse::exception::ParseError);
}

TEST(Parser, nargs_asterisk) {
  argparse::Parser *psr = new argparse::Parser("test");
  psr->add_argument("-a").nargs("*");
  psr->add_argument("-b");

  // Stop parsing by end of arguments.
  argparse::Argv ok1 = {"./test", "-a", "v1", "v2"};
  argparse::Values v1 = psr->parse_args(ok1);
  EXPECT_EQ(2, v1.size("a"));
  EXPECT_EQ("v1", v1.get("a", 0));
  EXPECT_EQ("v2", v1.get("a", 1));

  // Stop parsing by an other option.
  argparse::Argv ok2 = {"./test", "-a", "v1", "v2", "-b", "r1"};
  argparse::Values v2 = psr->parse_args(ok2);
  EXPECT_EQ(2, v2.size("a"));
  EXPECT_EQ("v1", v2.get("a", 0));
  EXPECT_EQ("v2", v2.get("a", 1));
  EXPECT_EQ(1, v2.size("b"));
  EXPECT_EQ("r1", v2.get("b"));

  // No option value.
  argparse::Argv ok3 = {"./test", "-a", "-b", "r1"};
  argparse::Values v3 = psr->parse_args(ok3);
  EXPECT_EQ(0, v3.size("a"));
  EXPECT_TRUE(v3.is_set("a"));
  EXPECT_EQ(1, v3.size("b"));
  EXPECT_EQ("r1", v3.get("b"));
}

class ParserNargsQuestion : public ::testing::Test {
public:
  argparse::Parser *psr;
  argparse::Argument *arg;
  virtual void SetUp() {
    psr = new argparse::Parser("test");
    arg = &(psr->add_argument("-a").nargs("?")); // 0 or 1 value
    psr->add_argument("-b");
  }

  virtual void TearDown() { delete psr; }
};

TEST_F(ParserNargsQuestion, ok0) {
  argparse::Argv seq = {"./test"};
  argparse::Values v = psr->parse_args(seq);
  EXPECT_FALSE(v.is_set("a"));
  EXPECT_EQ(0, v.size("a"));
  EXPECT_FALSE(v.is_set("b"));
}

TEST_F(ParserNargsQuestion, ok1) {
  argparse::Argv seq = {"./test", "-a"};
  argparse::Values v = psr->parse_args(seq);
  EXPECT_TRUE(v.is_set("a"));
  EXPECT_EQ(1, v.size("a"));
  EXPECT_FALSE(v.is_set("b"));
}

TEST_F(ParserNargsQuestion, ok2) {
  argparse::Argv seq = {"./test", "-a", "v1"};
  argparse::Values v = psr->parse_args(seq);
  EXPECT_TRUE(v.is_set("a"));
  EXPECT_EQ(1, v.size("a"));
  EXPECT_EQ("v1", v["a"]);
  EXPECT_FALSE(v.is_set("b"));
}

TEST_F(ParserNargsQuestion, ok3) {
  argparse::Argv seq = {"./test", "-a", "v1", "-b", "r1"};
  argparse::Values v = psr->parse_args(seq);
  EXPECT_TRUE(v.is_set("a"));
  EXPECT_EQ(1, v.size("a"));
  EXPECT_EQ("v1", v["a"]);
  EXPECT_TRUE(v.is_set("b"));
}

TEST_F(ParserNargsQuestion, with_default1) {
  arg->set_default("d");
  argparse::Argv seq = {"./test", "-a"};
  argparse::Values v = psr->parse_args(seq);
  EXPECT_TRUE(v.is_set("a"));
  EXPECT_EQ(1, v.size("a"));
}

TEST_F(ParserNargsQuestion, with_default2) {
  arg->set_default("d");
  argparse::Argv seq = {"./test"};
  argparse::Values v = psr->parse_args(seq);
  EXPECT_TRUE(v.is_set("a"));
  EXPECT_EQ(1, v.size("a"));
  EXPECT_EQ("d", v["a"]);
}

TEST_F(ParserNargsQuestion, with_const) {
  arg->set_const("c");
  argparse::Argv seq = {"./test", "-a"};
  argparse::Values v = psr->parse_args(seq);
  EXPECT_TRUE(v.is_set("a"));
  EXPECT_EQ(1, v.size("a"));
  EXPECT_EQ("c", v["a"]);
}

TEST_F(ParserNargsQuestion, with_default_and_const) {
  arg->set_const("c").set_default("d");
  argparse::Argv seq = {"./test", "-a"};
  argparse::Values v = psr->parse_args(seq);
  EXPECT_TRUE(v.is_set("a"));
  EXPECT_EQ(1, v.size("a"));
  // 'const' has priority over 'default'
  EXPECT_EQ("c", v["a"]);
}

class ParserNargsPlus : public ::testing::Test {
public:
  argparse::Parser *psr;
  argparse::Argument *arg;
  virtual void SetUp() {
    psr = new argparse::Parser("test");
    arg = &(psr->add_argument("-a").nargs("+")); // 0 or 1 value
    psr->add_argument("-b");
  }

  virtual void TearDown() { delete psr; }
};

TEST_F(ParserNargsPlus, ok0) {
  argparse::Argv seq = {"./test"};
  argparse::Values v = psr->parse_args(seq);
  EXPECT_FALSE(v.is_set("a"));
  EXPECT_EQ(0, v.size("a"));
  EXPECT_FALSE(v.is_set("b"));
}

TEST_F(ParserNargsPlus, ok1) {
  argparse::Argv seq = {"./test", "-a", "v1"};
  argparse::Values v = psr->parse_args(seq);
  EXPECT_TRUE(v.is_set("a"));
  EXPECT_EQ(1, v.size("a"));
  EXPECT_EQ("v1", v.get("a", 0));
  EXPECT_FALSE(v.is_set("b"));
}

TEST_F(ParserNargsPlus, ok2) {
  argparse::Argv seq = {"./test", "-a", "v1", "v2"};
  argparse::Values v = psr->parse_args(seq);
  EXPECT_TRUE(v.is_set("a"));
  EXPECT_EQ(2, v.size("a"));
  EXPECT_EQ("v1", v.get("a", 0));
  EXPECT_EQ("v2", v.get("a", 1));
  EXPECT_FALSE(v.is_set("b"));
}

TEST_F(ParserNargsPlus, ok3) {
  argparse::Argv seq = {"./test", "-a", "v1", "v2", "v3", "-b", "r1"};
  argparse::Values v = psr->parse_args(seq);
  EXPECT_TRUE(v.is_set("a"));
  EXPECT_EQ(3, v.size("a"));
  EXPECT_EQ("v1", v.get("a", 0));
  EXPECT_EQ("v2", v.get("a", 1));
  EXPECT_EQ("v3", v.get("a", 2));
  EXPECT_TRUE(v.is_set("b"));
}

TEST_F(ParserNargsPlus, ng_no_argument1) {
  argparse::Argv seq = {"./test", "-a"};
  EXPECT_THROW(psr->parse_args(seq), argparse::exception::ParseError);
}

TEST_F(ParserNargsPlus, ng_no_argument2) {
  argparse::Argv seq = {"./test", "-a", "-b", "r1"};
  EXPECT_THROW(psr->parse_args(seq), argparse::exception::ParseError);
}

TEST_F(ParserNargsPlus, with_default1) {
  arg->set_default("d");
  argparse::Argv seq = {"./test", "-a", "v1", "-b", "r1"};
  argparse::Values v = psr->parse_args(seq);
  EXPECT_EQ(1, v.size("a"));
  EXPECT_EQ("v1", v["a"]);
}

TEST_F(ParserNargsPlus, with_default2) {
  arg->set_default("d");
  argparse::Argv seq = {"./test", "-b", "r1"};
  argparse::Values v = psr->parse_args(seq);
  EXPECT_EQ(1, v.size("a"));
  EXPECT_EQ("d", v["a"]);
}

<commit_msg>use Fixture for Parser nargs number tests<commit_after>/*
 * Copyright 2016, Masayoshi Mizutani, mizutani@sfc.wide.ad.jp
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are
 * met:
 *
 * * Redistributions of source code must retain the above copyright
 * notice, this list of conditions and the following disclaimer.
 * * Redistributions in binary form must reproduce the above
 * copyright notice, this list of conditions and the following disclaimer
 * in the documentation and/or other materials provided with the
 * distribution.
 * * Neither the name of Google Inc. nor the names of its
 * contributors may be used to endorse or promote products derived from
 * this software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
 * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
 * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
 * A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
 * OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
 * SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
 * LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
 * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
 * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 */

#include <vector>
#include <string>

#include "./gtest.h"
#include "../argparse.hpp"

class ParserNargsNumber : public ::testing::Test {
public:
  argparse::Parser *psr;
  argparse::Argument *arg;
  virtual void SetUp() {
    psr = new argparse::Parser("test");
    arg = &(psr->add_argument("-a").nargs(2));
    psr->add_argument("-b");
  }

  virtual void TearDown() { delete psr; }
};

TEST_F(ParserNargsNumber, ok1) {
  argparse::Argv seq = {"./test", "-a", "v1", "v2"};
  argparse::Values val = psr->parse_args(seq);
  EXPECT_TRUE(val.is_set("a"));
  EXPECT_EQ(2, val.size("a"));
  EXPECT_EQ("v1", val.get("a", 0));
  EXPECT_EQ("v2", val.get("a", 1));
}

TEST_F(ParserNargsNumber, ng1_too_many_arg) {
  argparse::Argv seq = {"./test", "-a", "v1", "v2", "v3"};
  EXPECT_THROW(psr->parse_args(seq), argparse::exception::ParseError);
}

TEST_F(ParserNargsNumber, ng2_not_enough_arg) {
  argparse::Argv seq = {"./test", "-a", "v1"};
  EXPECT_THROW(psr->parse_args(seq), argparse::exception::ParseError);
}

class ParserNargsAsterisk : public ::testing::Test {
public:
  argparse::Parser *psr;
  argparse::Argument *arg;
  virtual void SetUp() {
    psr = new argparse::Parser("test");
    arg = &(psr->add_argument("-a").nargs("*"));
    psr->add_argument("-b");
  }

  virtual void TearDown() { delete psr; }
};

TEST_F(ParserNargsAsterisk, no_option) {
  argparse::Argv ok1 = {"./test"};
  argparse::Values v1 = psr->parse_args(ok1);
  EXPECT_FALSE(v1.is_set("a"));
}

TEST_F(ParserNargsAsterisk, ok1) {
  argparse::Argv ok1 = {"./test", "-a", "v1", "v2"};
  argparse::Values v1 = psr->parse_args(ok1);
  EXPECT_EQ(2, v1.size("a"));
  EXPECT_EQ("v1", v1.get("a", 0));
  EXPECT_EQ("v2", v1.get("a", 1));
}

TEST_F(ParserNargsAsterisk, ok2) {
  // Stop parsing by an other option.
  argparse::Argv ok2 = {"./test", "-a", "v1", "v2", "-b", "r1"};
  argparse::Values v2 = psr->parse_args(ok2);
  EXPECT_EQ(2, v2.size("a"));
  EXPECT_EQ("v1", v2.get("a", 0));
  EXPECT_EQ("v2", v2.get("a", 1));
  EXPECT_EQ(1, v2.size("b"));
  EXPECT_EQ("r1", v2.get("b"));
}

TEST_F(ParserNargsAsterisk, ok3) {
  // No option value.
  argparse::Argv ok3 = {"./test", "-a", "-b", "r1"};
  argparse::Values v3 = psr->parse_args(ok3);
  EXPECT_EQ(0, v3.size("a"));
  EXPECT_TRUE(v3.is_set("a"));
  EXPECT_EQ(1, v3.size("b"));
  EXPECT_EQ("r1", v3.get("b"));
}

class ParserNargsQuestion : public ::testing::Test {
public:
  argparse::Parser *psr;
  argparse::Argument *arg;
  virtual void SetUp() {
    psr = new argparse::Parser("test");
    arg = &(psr->add_argument("-a").nargs("?")); // 0 or 1 value
    psr->add_argument("-b");
  }

  virtual void TearDown() { delete psr; }
};

TEST_F(ParserNargsQuestion, ok0) {
  argparse::Argv seq = {"./test"};
  argparse::Values v = psr->parse_args(seq);
  EXPECT_FALSE(v.is_set("a"));
  EXPECT_EQ(0, v.size("a"));
  EXPECT_FALSE(v.is_set("b"));
}

TEST_F(ParserNargsQuestion, ok1) {
  argparse::Argv seq = {"./test", "-a"};
  argparse::Values v = psr->parse_args(seq);
  EXPECT_TRUE(v.is_set("a"));
  EXPECT_EQ(1, v.size("a"));
  EXPECT_FALSE(v.is_set("b"));
}

TEST_F(ParserNargsQuestion, ok2) {
  argparse::Argv seq = {"./test", "-a", "v1"};
  argparse::Values v = psr->parse_args(seq);
  EXPECT_TRUE(v.is_set("a"));
  EXPECT_EQ(1, v.size("a"));
  EXPECT_EQ("v1", v["a"]);
  EXPECT_FALSE(v.is_set("b"));
}

TEST_F(ParserNargsQuestion, ok3) {
  argparse::Argv seq = {"./test", "-a", "v1", "-b", "r1"};
  argparse::Values v = psr->parse_args(seq);
  EXPECT_TRUE(v.is_set("a"));
  EXPECT_EQ(1, v.size("a"));
  EXPECT_EQ("v1", v["a"]);
  EXPECT_TRUE(v.is_set("b"));
}

TEST_F(ParserNargsQuestion, with_default1) {
  arg->set_default("d");
  argparse::Argv seq = {"./test", "-a"};
  argparse::Values v = psr->parse_args(seq);
  EXPECT_TRUE(v.is_set("a"));
  EXPECT_EQ(1, v.size("a"));
}

TEST_F(ParserNargsQuestion, with_default2) {
  arg->set_default("d");
  argparse::Argv seq = {"./test"};
  argparse::Values v = psr->parse_args(seq);
  EXPECT_TRUE(v.is_set("a"));
  EXPECT_EQ(1, v.size("a"));
  EXPECT_EQ("d", v["a"]);
}

TEST_F(ParserNargsQuestion, with_const) {
  arg->set_const("c");
  argparse::Argv seq = {"./test", "-a"};
  argparse::Values v = psr->parse_args(seq);
  EXPECT_TRUE(v.is_set("a"));
  EXPECT_EQ(1, v.size("a"));
  EXPECT_EQ("c", v["a"]);
}

TEST_F(ParserNargsQuestion, with_default_and_const) {
  arg->set_const("c").set_default("d");
  argparse::Argv seq = {"./test", "-a"};
  argparse::Values v = psr->parse_args(seq);
  EXPECT_TRUE(v.is_set("a"));
  EXPECT_EQ(1, v.size("a"));
  // 'const' has priority over 'default'
  EXPECT_EQ("c", v["a"]);
}

class ParserNargsPlus : public ::testing::Test {
public:
  argparse::Parser *psr;
  argparse::Argument *arg;
  virtual void SetUp() {
    psr = new argparse::Parser("test");
    arg = &(psr->add_argument("-a").nargs("+")); // 0 or 1 value
    psr->add_argument("-b");
  }

  virtual void TearDown() { delete psr; }
};

TEST_F(ParserNargsPlus, ok0) {
  argparse::Argv seq = {"./test"};
  argparse::Values v = psr->parse_args(seq);
  EXPECT_FALSE(v.is_set("a"));
  EXPECT_EQ(0, v.size("a"));
  EXPECT_FALSE(v.is_set("b"));
}

TEST_F(ParserNargsPlus, ok1) {
  argparse::Argv seq = {"./test", "-a", "v1"};
  argparse::Values v = psr->parse_args(seq);
  EXPECT_TRUE(v.is_set("a"));
  EXPECT_EQ(1, v.size("a"));
  EXPECT_EQ("v1", v.get("a", 0));
  EXPECT_FALSE(v.is_set("b"));
}

TEST_F(ParserNargsPlus, ok2) {
  argparse::Argv seq = {"./test", "-a", "v1", "v2"};
  argparse::Values v = psr->parse_args(seq);
  EXPECT_TRUE(v.is_set("a"));
  EXPECT_EQ(2, v.size("a"));
  EXPECT_EQ("v1", v.get("a", 0));
  EXPECT_EQ("v2", v.get("a", 1));
  EXPECT_FALSE(v.is_set("b"));
}

TEST_F(ParserNargsPlus, ok3) {
  argparse::Argv seq = {"./test", "-a", "v1", "v2", "v3", "-b", "r1"};
  argparse::Values v = psr->parse_args(seq);
  EXPECT_TRUE(v.is_set("a"));
  EXPECT_EQ(3, v.size("a"));
  EXPECT_EQ("v1", v.get("a", 0));
  EXPECT_EQ("v2", v.get("a", 1));
  EXPECT_EQ("v3", v.get("a", 2));
  EXPECT_TRUE(v.is_set("b"));
}

TEST_F(ParserNargsPlus, ng_no_argument1) {
  argparse::Argv seq = {"./test", "-a"};
  EXPECT_THROW(psr->parse_args(seq), argparse::exception::ParseError);
}

TEST_F(ParserNargsPlus, ng_no_argument2) {
  argparse::Argv seq = {"./test", "-a", "-b", "r1"};
  EXPECT_THROW(psr->parse_args(seq), argparse::exception::ParseError);
}

TEST_F(ParserNargsPlus, with_default1) {
  arg->set_default("d");
  argparse::Argv seq = {"./test", "-a", "v1", "-b", "r1"};
  argparse::Values v = psr->parse_args(seq);
  EXPECT_EQ(1, v.size("a"));
  EXPECT_EQ("v1", v["a"]);
}

TEST_F(ParserNargsPlus, with_default2) {
  arg->set_default("d");
  argparse::Argv seq = {"./test", "-b", "r1"};
  argparse::Values v = psr->parse_args(seq);
  EXPECT_EQ(1, v.size("a"));
  EXPECT_EQ("d", v["a"]);
}

<|endoftext|>


<commit_before>/*
 * Copyright (C) 2004-2014 ZNC, see the NOTICE file for details.
 *
 * 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 <gtest/gtest.h>
#include <znc/ZNCString.h>

// GTest uses this function to output objects
void PrintTo(const CString& s, std::ostream* o) {
	*o << '"' << s.Escape_n(CString::EASCII, CString::EDEBUG) << '"';
}

class EscapeTest : public ::testing::Test {
protected:
	void testEncode(const CString& in, const CString& expectedOut, const CString& sformat) {
		CString::EEscape format = CString::ToEscape(sformat);
		CString out;

		SCOPED_TRACE("Format: " + sformat);

		// Encode, then decode again and check we still got the same string
		out = in.Escape_n(CString::EASCII, format);
		EXPECT_EQ(expectedOut, out);
		out = out.Escape_n(format, CString::EASCII);
		EXPECT_EQ(in, out);
	}

	void testString(const CString& in, const CString& url,
			const CString& html, const CString& sql) {
		SCOPED_TRACE("String: " + in);

		testEncode(in, url,  "URL");
		testEncode(in, html, "HTML");
		testEncode(in, sql,  "SQL");
	}
};

TEST_F(EscapeTest, Test) {
	//          input      url          html             sql
	testString("abcdefg", "abcdefg",   "abcdefg",       "abcdefg");
	testString("\n\t\r",  "%0A%09%0D", "\n\t\r",        "\\n\\t\\r");
	testString("'\"",     "%27%22",    "'"",       "\\'\\\"");
	testString("&<>",     "%26%3C%3E", "&<>", "&<>");
}

TEST(StringTest, Bool) {
	EXPECT_EQ(true, CString(true).ToBool());
	EXPECT_EQ(false, CString(false).ToBool());
}

#define CS(s) (CString((s), sizeof(s)-1))

TEST(StringTest, Cmp) {
	CString s = "Bbb";

	EXPECT_EQ(CString("Bbb"), s);
	EXPECT_LT(CString("Aaa"), s);
	EXPECT_GT(CString("Ccc"), s);
	EXPECT_EQ(0, s.StrCmp("Bbb"));
	EXPECT_GT(0, s.StrCmp("bbb"));
	EXPECT_LT(0, s.StrCmp("Aaa"));
	EXPECT_GT(0, s.StrCmp("Ccc"));
	EXPECT_EQ(0, s.CaseCmp("Bbb"));
	EXPECT_EQ(0, s.CaseCmp("bbb"));
	EXPECT_LT(0, s.CaseCmp("Aaa"));
	EXPECT_GT(0, s.CaseCmp("Ccc"));

	EXPECT_TRUE(s.Equals("bbb"));
	EXPECT_FALSE(s.Equals("bbb", true));
	EXPECT_FALSE(s.Equals("bb"));
}

TEST(StringTest, Wild) {
	EXPECT_TRUE(CString::WildCmp("*!?bar@foo", "I_am!~bar@foo"));
	EXPECT_TRUE(CString::WildCmp("", ""));
	EXPECT_TRUE(CString::WildCmp("*a*b*c*", "abc"));
	EXPECT_TRUE(CString::WildCmp("*a*b*c*", "axbyc"));
	EXPECT_FALSE(CString::WildCmp("*a*b*c*", "xy"));
}

TEST(StringTest, Case) {
	CString x = CS("xx");
	CString X = CS("XX");
	EXPECT_EQ(X, x.AsUpper());
	EXPECT_EQ(x, X.AsLower());
}

TEST(StringTest, Replace) {
	EXPECT_EQ("(b()b)", CString("(a()a)").Replace_n("a", "b"));
	EXPECT_EQ("(a()b)", CString("(a()a)").Replace_n("a", "b", "(", ")"));
	EXPECT_EQ("a(b)", CString("(a()a)").Replace_n("a", "b", "(", ")", true));
}

TEST(StringTest, Misc) {
	EXPECT_EQ("Hello,...", CString("Hello, I'm Bob").Ellipsize(9));
	EXPECT_EQ("Hello, I'm Bob", CString("Hello, I'm Bob").Ellipsize(90));
	EXPECT_EQ("..", CString("Hello, I'm Bob").Ellipsize(2));

	EXPECT_EQ("Xy", CS("Xyz").Left(2));
	EXPECT_EQ("Xyz", CS("Xyz").Left(20));

	EXPECT_EQ("yz", CS("Xyz").Right(2));
	EXPECT_EQ("Xyz", CS("Xyz").Right(20));
}

TEST(StringTest, Split) {
	EXPECT_EQ("a", CS("a b c").Token(0));
	EXPECT_EQ("b", CS("a b c").Token(1));
	EXPECT_EQ("", CS("a b c").Token(100));
	EXPECT_EQ("b c", CS("a b c").Token(1, true));
	EXPECT_EQ("c", CS("a  c").Token(1));
	EXPECT_EQ("", CS("a  c").Token(1, false, " ", true));
	EXPECT_EQ("c", CS("a  c").Token(1, false, "  "));
	EXPECT_EQ(" c", CS("a   c").Token(1, false, "  "));
	EXPECT_EQ("c", CS("a    c").Token(1, false, "  "));
	EXPECT_EQ("b c", CS("a (b c) d").Token(1, false, " ", false, "(", ")"));
	EXPECT_EQ("(b c)", CS("a (b c) d").Token(1, false, " ", false, "(", ")", false));
	EXPECT_EQ("d", CS("a (b c) d").Token(2, false, " ", false, "(", ")", false));

	VCString vexpected;
	VCString vresult;

	vexpected.push_back("a");
	vexpected.push_back("b");
	vexpected.push_back("c");
	CS("a b c").Split(" ", vresult);
	EXPECT_EQ(vexpected, vresult);

	MCString mexpected;
	MCString mresult;

	mexpected["a"] = "b";
	mexpected["c"] = "d";
	CS("a=x&c=d&a=b").URLSplit(mresult);
	EXPECT_EQ(mexpected, mresult) << "URLSplit";
}

TEST(StringTest, NamedFormat) {
	MCString m;
	m["a"] = "b";
	EXPECT_EQ("{xbyb", CString::NamedFormat(CS("\\{x{a}y{a}"), m));
}

TEST(StringTest, Hash) {
	EXPECT_EQ("d41d8cd98f00b204e9800998ecf8427e", CS("").MD5());
	EXPECT_EQ("0cc175b9c0f1b6a831c399e269772661", CS("a").MD5());

	EXPECT_EQ("e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855", CS("").SHA256());
	EXPECT_EQ("ca978112ca1bbdcafac231b39a23dc4da786eff8147c4e72b9807785afee48bb", CS("a").SHA256());
}
<commit_msg>StringTest: Make a local function static<commit_after>/*
 * Copyright (C) 2004-2014 ZNC, see the NOTICE file for details.
 *
 * 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 <gtest/gtest.h>
#include <znc/ZNCString.h>

// GTest uses this function to output objects
static void PrintTo(const CString& s, std::ostream* o) {
	*o << '"' << s.Escape_n(CString::EASCII, CString::EDEBUG) << '"';
}

class EscapeTest : public ::testing::Test {
protected:
	void testEncode(const CString& in, const CString& expectedOut, const CString& sformat) {
		CString::EEscape format = CString::ToEscape(sformat);
		CString out;

		SCOPED_TRACE("Format: " + sformat);

		// Encode, then decode again and check we still got the same string
		out = in.Escape_n(CString::EASCII, format);
		EXPECT_EQ(expectedOut, out);
		out = out.Escape_n(format, CString::EASCII);
		EXPECT_EQ(in, out);
	}

	void testString(const CString& in, const CString& url,
			const CString& html, const CString& sql) {
		SCOPED_TRACE("String: " + in);

		testEncode(in, url,  "URL");
		testEncode(in, html, "HTML");
		testEncode(in, sql,  "SQL");
	}
};

TEST_F(EscapeTest, Test) {
	//          input      url          html             sql
	testString("abcdefg", "abcdefg",   "abcdefg",       "abcdefg");
	testString("\n\t\r",  "%0A%09%0D", "\n\t\r",        "\\n\\t\\r");
	testString("'\"",     "%27%22",    "'"",       "\\'\\\"");
	testString("&<>",     "%26%3C%3E", "&<>", "&<>");
}

TEST(StringTest, Bool) {
	EXPECT_EQ(true, CString(true).ToBool());
	EXPECT_EQ(false, CString(false).ToBool());
}

#define CS(s) (CString((s), sizeof(s)-1))

TEST(StringTest, Cmp) {
	CString s = "Bbb";

	EXPECT_EQ(CString("Bbb"), s);
	EXPECT_LT(CString("Aaa"), s);
	EXPECT_GT(CString("Ccc"), s);
	EXPECT_EQ(0, s.StrCmp("Bbb"));
	EXPECT_GT(0, s.StrCmp("bbb"));
	EXPECT_LT(0, s.StrCmp("Aaa"));
	EXPECT_GT(0, s.StrCmp("Ccc"));
	EXPECT_EQ(0, s.CaseCmp("Bbb"));
	EXPECT_EQ(0, s.CaseCmp("bbb"));
	EXPECT_LT(0, s.CaseCmp("Aaa"));
	EXPECT_GT(0, s.CaseCmp("Ccc"));

	EXPECT_TRUE(s.Equals("bbb"));
	EXPECT_FALSE(s.Equals("bbb", true));
	EXPECT_FALSE(s.Equals("bb"));
}

TEST(StringTest, Wild) {
	EXPECT_TRUE(CString::WildCmp("*!?bar@foo", "I_am!~bar@foo"));
	EXPECT_TRUE(CString::WildCmp("", ""));
	EXPECT_TRUE(CString::WildCmp("*a*b*c*", "abc"));
	EXPECT_TRUE(CString::WildCmp("*a*b*c*", "axbyc"));
	EXPECT_FALSE(CString::WildCmp("*a*b*c*", "xy"));
}

TEST(StringTest, Case) {
	CString x = CS("xx");
	CString X = CS("XX");
	EXPECT_EQ(X, x.AsUpper());
	EXPECT_EQ(x, X.AsLower());
}

TEST(StringTest, Replace) {
	EXPECT_EQ("(b()b)", CString("(a()a)").Replace_n("a", "b"));
	EXPECT_EQ("(a()b)", CString("(a()a)").Replace_n("a", "b", "(", ")"));
	EXPECT_EQ("a(b)", CString("(a()a)").Replace_n("a", "b", "(", ")", true));
}

TEST(StringTest, Misc) {
	EXPECT_EQ("Hello,...", CString("Hello, I'm Bob").Ellipsize(9));
	EXPECT_EQ("Hello, I'm Bob", CString("Hello, I'm Bob").Ellipsize(90));
	EXPECT_EQ("..", CString("Hello, I'm Bob").Ellipsize(2));

	EXPECT_EQ("Xy", CS("Xyz").Left(2));
	EXPECT_EQ("Xyz", CS("Xyz").Left(20));

	EXPECT_EQ("yz", CS("Xyz").Right(2));
	EXPECT_EQ("Xyz", CS("Xyz").Right(20));
}

TEST(StringTest, Split) {
	EXPECT_EQ("a", CS("a b c").Token(0));
	EXPECT_EQ("b", CS("a b c").Token(1));
	EXPECT_EQ("", CS("a b c").Token(100));
	EXPECT_EQ("b c", CS("a b c").Token(1, true));
	EXPECT_EQ("c", CS("a  c").Token(1));
	EXPECT_EQ("", CS("a  c").Token(1, false, " ", true));
	EXPECT_EQ("c", CS("a  c").Token(1, false, "  "));
	EXPECT_EQ(" c", CS("a   c").Token(1, false, "  "));
	EXPECT_EQ("c", CS("a    c").Token(1, false, "  "));
	EXPECT_EQ("b c", CS("a (b c) d").Token(1, false, " ", false, "(", ")"));
	EXPECT_EQ("(b c)", CS("a (b c) d").Token(1, false, " ", false, "(", ")", false));
	EXPECT_EQ("d", CS("a (b c) d").Token(2, false, " ", false, "(", ")", false));

	VCString vexpected;
	VCString vresult;

	vexpected.push_back("a");
	vexpected.push_back("b");
	vexpected.push_back("c");
	CS("a b c").Split(" ", vresult);
	EXPECT_EQ(vexpected, vresult);

	MCString mexpected;
	MCString mresult;

	mexpected["a"] = "b";
	mexpected["c"] = "d";
	CS("a=x&c=d&a=b").URLSplit(mresult);
	EXPECT_EQ(mexpected, mresult) << "URLSplit";
}

TEST(StringTest, NamedFormat) {
	MCString m;
	m["a"] = "b";
	EXPECT_EQ("{xbyb", CString::NamedFormat(CS("\\{x{a}y{a}"), m));
}

TEST(StringTest, Hash) {
	EXPECT_EQ("d41d8cd98f00b204e9800998ecf8427e", CS("").MD5());
	EXPECT_EQ("0cc175b9c0f1b6a831c399e269772661", CS("a").MD5());

	EXPECT_EQ("e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855", CS("").SHA256());
	EXPECT_EQ("ca978112ca1bbdcafac231b39a23dc4da786eff8147c4e72b9807785afee48bb", CS("a").SHA256());
}
<|endoftext|>


<commit_before>#include <pex/PexFunctionBuilder.h>

namespace caprica { namespace pex {

void PexFunctionBuilder::populateFunction(PexFunction* func, PexDebugFunctionInfo* debInfo) {
  for (size_t i = 0; i < instructions.size(); i++) {
    for (auto& arg : instructions[i]->args) {
      if (arg.type == PexValueType::Label) {
        if (arg.l->targetIdx == (size_t)-1)
          CapricaError::logicalFatal("Unresolved label!");
        auto newVal = arg.l->targetIdx - i;
        arg.type = PexValueType::Integer;
        arg.i = (int32_t)newVal;
      }
    }
  }

  for (auto l : labels) {
    if (l->targetIdx == (size_t)-1)
      CapricaError::logicalFatal("Unused unresolved label!");
    delete l;
  }
  labels.clear();

  for (auto tmp : tempVarRefs) {
    if (tmp->var == nullptr)
      CapricaError::logicalFatal("Unresolved tmp var!");
    delete tmp;
  }
  tempVarRefs.clear();

  func->instructions = instructions;
  func->locals = locals;
  debInfo->instructionLineMap.reserve(instructionLocations.size());
  for (auto& l : instructionLocations) {
    if (l.line > std::numeric_limits<uint16_t>::max())
      CapricaError::fatal(l, "The file has too many lines for the debug info to be able to map correctly!");
    debInfo->instructionLineMap.push_back((uint16_t)l.line);
  }
}

static int32_t getDestArgIndexForOpCode(PexOpCode op) {
  switch (op) {
    case PexOpCode::Nop:
      return -1;
    case PexOpCode::CallMethod:
    case PexOpCode::CallStatic:
      return 2;
    case PexOpCode::CallParent:
      return 1;
#define OP_ARG1(name, opcode, destArgIdx, at1, an1) \
    case PexOpCode::opcode: return destArgIdx;
#define OP_ARG2(name, opcode, destArgIdx, at1, an1, at2, an2) \
    case PexOpCode::opcode: return destArgIdx;
#define OP_ARG3(name, opcode, destArgIdx, at1, an1, at2, an2, at3, an3) \
    case PexOpCode::opcode: return destArgIdx;
#define OP_ARG4(name, opcode, destArgIdx, at1, an1, at2, an2, at3, an3, at4, an4) \
    case PexOpCode::opcode: return destArgIdx;
#define OP_ARG5(name, opcode, destArgIdx, at1, an1, at2, an2, at3, an3, at4, an4, at5, an5) \
    case PexOpCode::opcode: return destArgIdx;
OPCODES(OP_ARG1, OP_ARG2, OP_ARG3, OP_ARG4, OP_ARG5)
#undef OP_ARG1
#undef OP_ARG2
#undef OP_ARG3
#undef OP_ARG4
#undef OP_ARG5
    default:
      CapricaError::logicalFatal("Unknown PexOpCode!");
  }
}

void PexFunctionBuilder::freeValueIfTemp(const PexValue& v) {
  PexString varName;
  if (v.type == PexValueType::Identifier)
    varName = v.s;
  else if (v.type == PexValueType::TemporaryVar && v.tmpVar->var)
    varName = v.tmpVar->var->name;
  else
    return;

  auto f = tempVarNameTypeMap.find(varName);
  if (f != tempVarNameTypeMap.end() && !longLivedTempVars.count(varName)) {
    if (!freeTempVars.count(f->second->type))
      freeTempVars.insert({ f->second->type, { } });
    freeTempVars[f->second->type].push_back(f->second);
  }
}

PexLocalVariable* PexFunctionBuilder::internalAllocateTempVar(const PexString& typeName) {
  auto f = freeTempVars.find(typeName);
  if (f != freeTempVars.end() && f->second.size()) {
    PexLocalVariable* b = f->second.back();
    f->second.pop_back();
    return b;
  }

  auto loc = new PexLocalVariable();
  std::stringstream ss;
  ss << "::temp" << currentTempI++;
  loc->name = file->getString(ss.str());
  loc->type = typeName;
  tempVarNameTypeMap.insert({ loc->name, loc });
  locals.push_back(loc);
  return loc;
}

PexFunctionBuilder& PexFunctionBuilder::push(PexInstruction* instr) {
  for (auto& v : instr->args) {
    if (v.type == PexValueType::Invalid)
      CapricaError::fatal(currentLocation, "Attempted to use an invalid value as a value! (perhaps you tried to use the return value of a function that doesn't return?)");
    if (v.type == PexValueType::TemporaryVar && v.tmpVar->var)
      v = PexValue::Identifier(v.tmpVar->var);
    freeValueIfTemp(v);
  }
  for (auto& v : instr->variadicArgs) {
    if (v.type == PexValueType::Invalid)
      CapricaError::fatal(currentLocation, "Attempted to use an invalid value as a value! (perhaps you tried to use the return value of a function that doesn't return?)");
    if (v.type == PexValueType::TemporaryVar && v.tmpVar->var)
      v = PexValue::Identifier(v.tmpVar->var);
    freeValueIfTemp(v);
  }

  auto destIdx = getDestArgIndexForOpCode(instr->opCode);
  if (destIdx != -1 && instr->args[destIdx].type == PexValueType::TemporaryVar) {
    auto loc = internalAllocateTempVar(instr->args[destIdx].tmpVar->type);
    instr->args[destIdx].tmpVar->var = loc;
    instr->args[destIdx] = PexValue::Identifier(loc);
  }

  for (auto& v : instr->args) {
    if (v.type == PexValueType::TemporaryVar)
      CapricaError::fatal(currentLocation, "Attempted to use a temporary var before it's been assigned!");
  }
  for (auto& v : instr->variadicArgs) {
    if (v.type == PexValueType::TemporaryVar)
      CapricaError::fatal(currentLocation, "Attempted to use a temporary var before it's been assigned!");
  }

  // Special case for dead assigns.
  if (CapricaConfig::enableOptimizations && instr->opCode == PexOpCode::Assign) {
    auto dst = instr->args[0];
    auto src = instr->args[1];
    if (dst.type == src.type && dst.type == PexValueType::Identifier && dst.s.index == src.s.index)
      return *this;
  }

  instructionLocations.push_back(currentLocation);
  instructions.push_back(instr);
  return *this;
}

}}
<commit_msg>More cleanup.<commit_after>#include <pex/PexFunctionBuilder.h>

namespace caprica { namespace pex {

void PexFunctionBuilder::populateFunction(PexFunction* func, PexDebugFunctionInfo* debInfo) {
  for (size_t i = 0; i < instructions.size(); i++) {
    for (auto& arg : instructions[i]->args) {
      if (arg.type == PexValueType::Label) {
        if (arg.l->targetIdx == (size_t)-1)
          CapricaError::logicalFatal("Unresolved label!");
        auto newVal = arg.l->targetIdx - i;
        arg.type = PexValueType::Integer;
        arg.i = (int32_t)newVal;
      }
    }
  }

  for (auto l : labels) {
    if (l->targetIdx == (size_t)-1)
      CapricaError::logicalFatal("Unused unresolved label!");
    delete l;
  }
  labels.clear();

  for (auto tmp : tempVarRefs) {
    if (tmp->var == nullptr)
      CapricaError::logicalFatal("Unresolved tmp var!");
    delete tmp;
  }
  tempVarRefs.clear();

  func->instructions = instructions;
  func->locals = locals;
  debInfo->instructionLineMap.reserve(instructionLocations.size());
  for (auto& l : instructionLocations) {
    if (l.line > std::numeric_limits<uint16_t>::max())
      CapricaError::fatal(l, "The file has too many lines for the debug info to be able to map correctly!");
    debInfo->instructionLineMap.push_back((uint16_t)l.line);
  }
}

static int32_t getDestArgIndexForOpCode(PexOpCode op) {
  switch (op) {
    case PexOpCode::Nop:
      return -1;
    case PexOpCode::CallMethod:
    case PexOpCode::CallStatic:
      return 2;
    case PexOpCode::CallParent:
      return 1;
#define OP_ARG1(name, opcode, destArgIdx, ...) case PexOpCode::opcode: return destArgIdx;
#define OP_ARG2(name, opcode, destArgIdx, ...) case PexOpCode::opcode: return destArgIdx;
#define OP_ARG3(name, opcode, destArgIdx, ...) case PexOpCode::opcode: return destArgIdx;
#define OP_ARG4(name, opcode, destArgIdx, ...) case PexOpCode::opcode: return destArgIdx;
#define OP_ARG5(name, opcode, destArgIdx, ...) case PexOpCode::opcode: return destArgIdx;
OPCODES(OP_ARG1, OP_ARG2, OP_ARG3, OP_ARG4, OP_ARG5)
#undef OP_ARG1
#undef OP_ARG2
#undef OP_ARG3
#undef OP_ARG4
#undef OP_ARG5
    default:
      CapricaError::logicalFatal("Unknown PexOpCode!");
  }
}

void PexFunctionBuilder::freeValueIfTemp(const PexValue& v) {
  PexString varName;
  if (v.type == PexValueType::Identifier)
    varName = v.s;
  else if (v.type == PexValueType::TemporaryVar && v.tmpVar->var)
    varName = v.tmpVar->var->name;
  else
    return;

  auto f = tempVarNameTypeMap.find(varName);
  if (f != tempVarNameTypeMap.end() && !longLivedTempVars.count(varName)) {
    if (!freeTempVars.count(f->second->type))
      freeTempVars.insert({ f->second->type, { } });
    freeTempVars[f->second->type].push_back(f->second);
  }
}

PexLocalVariable* PexFunctionBuilder::internalAllocateTempVar(const PexString& typeName) {
  auto f = freeTempVars.find(typeName);
  if (f != freeTempVars.end() && f->second.size()) {
    PexLocalVariable* b = f->second.back();
    f->second.pop_back();
    return b;
  }

  auto loc = new PexLocalVariable();
  std::stringstream ss;
  ss << "::temp" << currentTempI++;
  loc->name = file->getString(ss.str());
  loc->type = typeName;
  tempVarNameTypeMap.insert({ loc->name, loc });
  locals.push_back(loc);
  return loc;
}

PexFunctionBuilder& PexFunctionBuilder::push(PexInstruction* instr) {
  for (auto& v : instr->args) {
    if (v.type == PexValueType::Invalid)
      CapricaError::fatal(currentLocation, "Attempted to use an invalid value as a value! (perhaps you tried to use the return value of a function that doesn't return?)");
    if (v.type == PexValueType::TemporaryVar && v.tmpVar->var)
      v = PexValue::Identifier(v.tmpVar->var);
    freeValueIfTemp(v);
  }
  for (auto& v : instr->variadicArgs) {
    if (v.type == PexValueType::Invalid)
      CapricaError::fatal(currentLocation, "Attempted to use an invalid value as a value! (perhaps you tried to use the return value of a function that doesn't return?)");
    if (v.type == PexValueType::TemporaryVar && v.tmpVar->var)
      v = PexValue::Identifier(v.tmpVar->var);
    freeValueIfTemp(v);
  }

  auto destIdx = getDestArgIndexForOpCode(instr->opCode);
  if (destIdx != -1 && instr->args[destIdx].type == PexValueType::TemporaryVar) {
    auto loc = internalAllocateTempVar(instr->args[destIdx].tmpVar->type);
    instr->args[destIdx].tmpVar->var = loc;
    instr->args[destIdx] = PexValue::Identifier(loc);
  }

  for (auto& v : instr->args) {
    if (v.type == PexValueType::TemporaryVar)
      CapricaError::fatal(currentLocation, "Attempted to use a temporary var before it's been assigned!");
  }
  for (auto& v : instr->variadicArgs) {
    if (v.type == PexValueType::TemporaryVar)
      CapricaError::fatal(currentLocation, "Attempted to use a temporary var before it's been assigned!");
  }

  // Special case for dead assigns.
  if (CapricaConfig::enableOptimizations && instr->opCode == PexOpCode::Assign) {
    auto dst = instr->args[0];
    auto src = instr->args[1];
    if (dst.type == src.type && dst.type == PexValueType::Identifier && dst.s.index == src.s.index)
      return *this;
  }

  instructionLocations.push_back(currentLocation);
  instructions.push_back(instr);
  return *this;
}

}}
<|endoftext|>