<commit_before>#include "sift.hpp"

#include <iostream>
#include <string>
#include <cassert>

#include <vigra/impex.hxx>
#include <vigra/multi_math.hxx>
#include <vigra/linear_algebra.hxx>

#include "point.hpp"
#include "algorithms.hpp"

using namespace vigra::multi_math;
using namespace vigra::linalg;

namespace sift {
    void Sift::calculate(vigra::MultiArray<2, f32_t>& img) {
        auto dogs = _createDOGs(img);
        //Save DoGs for Demonstration purposes
        for (u16_t i = 0; i < dogs.width(); i++) {
            for (u16_t j = 0; j < dogs.height(); j++) {
                const std::string fnStr = "images/dog" + std::to_string(i) + std::to_string(j) + ".png";
                exportImage(dogs(i, j), vigra::ImageExportInfo(fnStr.c_str()));
            }
        }

        auto interestPoints = _findScaleSpaceExtrema(dogs);
        //Save img with found interest points for demonstration purposes
        auto img_output1 = img;
        for (u16_t i = 0; i < interestPoints.height(); i++) {
            for (u16_t x = 0; x < img.width(); x++) {
                for (u16_t y = 0; y < img.height(); y++) {
                    if (interestPoints(0, i)(x, y) > -1) {
                        img_output1(x, y) = 255;
                    }
                }
            }
        }

        exportImage(img_output1, vigra::ImageExportInfo("images/interest_points.png"));
        _eliminateEdgeResponses(interestPoints, dogs);
        //Save img with filtered interest points for demonstration purposes
        auto img_output2 = img;
        for (u16_t i = 0; i < interestPoints.height(); i++) {
            for (u16_t x = 0; x < img.width(); x++) {
                for (u16_t y = 0; y < img.height(); y++) {
                    if (interestPoints(0, i)(x, y) > -1) {
                        img_output2(x, y) = 255;
                    }
                }
            }
        }
        exportImage(img_output2, vigra::ImageExportInfo("images/after_filter.png"));
        _orientationAssignment(interestPoints);
    }

    void Sift::_orientationAssignment(Matrix<Matrix<f32_t>> interestPoints) {
        for (u16_t e = 0; e < 1; e++) {
            for (u16_t i = 0; i < interestPoints.height(); i++) {
                auto closest = _gaussians(e, i);
                f32_t scale = _calculateScale(e, i);
                vigra::Shape2 size(closest.width(), closest.height());
                vigra::MultiArray<2, f32_t> magnitudes(size);
                vigra::MultiArray<2, f32_t> orientations(size);
                for (u16_t x = 0; x < 1; x++) {
                    for (u16_t y = 0; y < closest.height(); y++) {
                        Point p(x, y);
                        magnitudes(x, y) = alg::gradientMagnitude(closest, p);
                        orientations(x, y) = alg::gradientOrientation(closest, p);
                    }
                }
                for (u16_t x = 8; x < interestPoints(e, i).width() - 8; x++) {
                    for (u16_t y = 8; y < interestPoints(e, i).height() - 8; y++) {
                        if (interestPoints(e, i)(x, y) >= 0) {
                            auto topLeftCorner = vigra::Shape2(x - 8, y - 8);
                            auto bottomRightCorner = vigra::Shape2(x + 8, y + 8);

                            auto orientation_region = orientations.subarray(topLeftCorner, bottomRightCorner);
                            auto mag_region = magnitudes.subarray(topLeftCorner, bottomRightCorner);

                            auto histogram = alg::orientationHistogram(orientation_region, mag_region, scale);

                            _createPeak(histogram);

                            //TODO:: Parabola to the 3 values closest to the peak

                        }
                    }
                }
            }
        }
    }

    void Sift::_createPeak(std::array<f32_t, 36>& histo) {
        const f32_t max = *(std::max_element(histo.begin(), histo.end()));
        //allowed range(80% of max)
        const f32_t range = max / 5 * 4;
        std::for_each(histo.begin(), histo.end(), [&](f32_t& elem) { if (elem < range) elem = -1; });
    }

    f32_t Sift::_calculateScale(u16_t e, u16_t i) const {
        if (e < 1) {
            if (i < 1)  {
                return _sigma;
            }
            return std::pow(_k, i) * _sigma;
        } else if (e == 1) {
            return std::pow(_k, _dogsPerEpoch + i) * _sigma;
        }
        //_dogsPerEpoch + 2: count of gaussians per epoch
        //e - 1: the very last epoch before our
        // - 2: take the top img from the stack and add i to it
        return std::pow(_k, (_dogsPerEpoch + 2) * (e - 1) - 2 + i) * _sigma;
    }

    void Sift::_eliminateEdgeResponses(Matrix<Matrix<f32_t>>& interestPoints,
            const Matrix<vigra::MultiArray<2, f32_t>>& dogs) const {

        for(u16_t e = 0; e < dogs.width(); e++) {
            for (u16_t i = 1; i < dogs.height() - 1; i++) {
                for (u16_t x = 1; x < dogs(e, i).shape(0) - 1; x++) {
                    for (u16_t y = 1; y < dogs(e, i).shape(1) - 1; y++) {
                        if (interestPoints(e, i - 1)(x, y) > -1) {
                            auto d = dogs(e, i);

                            const vigra::MultiArray<2, f32_t> param[3] = {dogs(e, i - 1), dogs(e, i), dogs(e, i + 1)};
                            const Point p(x, y);
                            vigra::Matrix<f32_t> deriv = alg::foDerivative(param, p);
                            vigra::Matrix<f32_t> sec_deriv = alg::soDerivative(param, p);

                            vigra::Matrix<f32_t> neg_sec_deriv = sec_deriv ;
                            neg_sec_deriv *=  -1;

                            vigra::MultiArray<2, f32_t> extremum(vigra::Shape2(3, 1));
                            if (!linearSolve(inverse(neg_sec_deriv), deriv, extremum)) {
                                std::cerr << "Couldn't solve linear system" << std::endl;
                                throw;
                            }

                            //Calculated up 0.5 from paper to own image values [0,255]
                            if (extremum(0, 0) > 127.5 || extremum(1, 0) > 127.5 || extremum(2, 0) > 127.5) {
                                interestPoints(e, i - 1)(x, y) = -1;
                                continue;
                            }
                            vigra::Matrix<f32_t> deriv_transpose = deriv.transpose();
                            f32_t func_val_extremum = dot(deriv_transpose, extremum);
                            func_val_extremum *= 0.5 + d(x, y);

                            //Calculated up 0.03 from paper to own image values[0, 255]
                            if (func_val_extremum < 7.65) {
                                interestPoints(e, i - 1)(x, y) = -1;
                                continue;
                            }

                            //dxx + dyy
                            f32_t hessian_tr = sec_deriv(0, 0) + sec_deriv(1, 1);
                            //dxx * dyy - dxy^2
                            f32_t hessian_det = sec_deriv(0, 0) * sec_deriv(1, 1) - std::pow(sec_deriv(0, 1), 2);

                            if (hessian_det < 0) {
                                interestPoints(e, i - 1)(x, y) = -1;
                                continue;
                            }

                            //Original r = 10, calculated up to own image values[0, 255]
                            if (std::pow(hessian_tr, 2) / hessian_det > std::pow(2550 + 1, 2) / 2550) {
                                interestPoints(e, i - 1)(x, y) = -1;
                            }
                        }
                    }
                }
            }
        }
    }

    const Matrix<Matrix<f32_t>> Sift::_findScaleSpaceExtrema(const Matrix<vigra::MultiArray<2, f32_t>>& dogs) const {
        //A matrix of matrix. Outer dogs will be ignored, because we need a upper and lower neighbor
        Matrix<Matrix<f32_t>> interestPoints(dogs.width(), dogs.height() - 2);
        for (u16_t e = 0; e < dogs.width(); e++) {
            for (u16_t i = 1; i < dogs.height() - 1; i++) {
                interestPoints(e, i - 1) = Matrix<f32_t>(dogs(e, i).width(), dogs(e, i).height(), -1);
                for (i16_t x = 0; x < dogs(e, i).shape(0); x++) {
                    for (i16_t y = 0; y < dogs(e, i).shape(1); y++) {
                        auto leftUpCorner = vigra::Shape2(x - 1, y - 1);
                        auto rightDownCorner = vigra::Shape2(x + 1, y + 1);

                        //Get the neighborhood of the current pixel
                        auto current = dogs(e, i).subarray(leftUpCorner, rightDownCorner);
                        //Get neighborhood of adjacent DOGs
                        auto under = dogs(e, i - 1).subarray(leftUpCorner, rightDownCorner);
                        auto above = dogs(e, i + 1).subarray(leftUpCorner, rightDownCorner);
                        //Check all neighborhood pixels of current and adjacent DOGs. If there isn't any
                        //pixel bigger or smaller than the current. We found an extremum.
                        if ((!any(current > dogs(e, i)(x, y)) &&
                                    !any(under > dogs(e, i)(x, y)) &&
                                    !any(above > dogs(e, i)(x, y))) ||
                                (!any(current < dogs(e, i)(x, y)) &&
                                 !any(under < dogs(e, i)(x, y)) &&
                                 !any(above < dogs(e, i)(x, y))))
                        {
                            interestPoints(e, i - 1)(x, y) = dogs(e, i)(x, y);
                        }
                    }
                }
            }
        }
        return interestPoints; // TODO: by ref entgegen nehmen, um copy zu vermeiden?
    }

    const Matrix<vigra::MultiArray<2, f32_t>> Sift::_createDOGs(vigra::MultiArray<2, f32_t>& img) {

        assert(_epochs > 0); // pre condition
        assert(_dogsPerEpoch >= 3); // pre condition

        Matrix<vigra::MultiArray<2, f32_t>> gaussians(_epochs, _dogsPerEpoch + 2);
        Matrix<vigra::MultiArray<2, f32_t>> dogs(_epochs, _dogsPerEpoch);

        gaussians(0, 0) = alg::convolveWithGauss(img, _sigma);

        //TODO: More elegant way?
        u16_t exp = 0;
        for (i16_t i = 0; i < _epochs; i++) {
            for (i16_t j = 1; j < _dogsPerEpoch + 1; j++) {
                gaussians(i, j) = alg::convolveWithGauss(gaussians(i, j - 1), std::pow(_k, exp) * _sigma);
                dogs(i, j - 1) = alg::dog(gaussians(i, j - 1), gaussians(i, j));
                exp++;
            }
            /*
             * If we aren't in the last epoch populate the next level with the second
             * last element, scaled by a half, of images of current epoch.
             */
            if (i < (_epochs - 1)) {
                auto scaledElem = alg::reduceToNextLevel(gaussians(i, _dogsPerEpoch - 1), _sigma);
                gaussians(i + 1, 0) = scaledElem;

                exp -= 2;
            }
        }
        _gaussians = gaussians;
        return dogs; // TODO: by ref entgegen nehmen um copy zu vermeiden?
    }
}
<commit_msg>offset at hessian set to 10 and some minor fixes<commit_after>#include "sift.hpp"

#include <iostream>
#include <string>
#include <cassert>

#include <vigra/impex.hxx>
#include <vigra/multi_math.hxx>
#include <vigra/linear_algebra.hxx>

#include "point.hpp"
#include "algorithms.hpp"

using namespace vigra::multi_math;
using namespace vigra::linalg;

namespace sift {
    void Sift::calculate(vigra::MultiArray<2, f32_t>& img) {
        auto dogs = _createDOGs(img);
        //Save DoGs for Demonstration purposes
        for (u16_t i = 0; i < dogs.width(); i++) {
            for (u16_t j = 0; j < dogs.height(); j++) {
                const std::string fnStr = "images/dog" + std::to_string(i) + std::to_string(j) + ".png";
                exportImage(dogs(i, j), vigra::ImageExportInfo(fnStr.c_str()));
            }
        }

        auto interestPoints = _findScaleSpaceExtrema(dogs);
        //Save img with found interest points for demonstration purposes
        auto img_output1 = img;
        for (u16_t i = 0; i < interestPoints.height(); i++) {
            for (u16_t x = 0; x < img.width(); x++) {
                for (u16_t y = 0; y < img.height(); y++) {
                    if (interestPoints(0, i)(x, y) > -1) {
                        img_output1(x, y) = 255;
                    }
                }
            }
        }

        exportImage(img_output1, vigra::ImageExportInfo("images/interest_points.png"));
        _eliminateEdgeResponses(interestPoints, dogs);
        //Save img with filtered interest points for demonstration purposes
        auto img_output2 = img;
        for (u16_t i = 0; i < interestPoints.height(); i++) {
            for (u16_t x = 0; x < img.width(); x++) {
                for (u16_t y = 0; y < img.height(); y++) {
                    if (interestPoints(0, i)(x, y) > -1) {
                        img_output2(x, y) = 255;
                    }
                }
            }
        }
        exportImage(img_output2, vigra::ImageExportInfo("images/after_filter.png"));
        _orientationAssignment(interestPoints);
    }

    void Sift::_orientationAssignment(Matrix<Matrix<f32_t>> interestPoints) {
        for (u16_t e = 0; e < 1; e++) {
            for (u16_t i = 0; i < interestPoints.height(); i++) {
                auto closest = _gaussians(e, i);
                f32_t scale = _calculateScale(e, i);
                vigra::Shape2 size(closest.width(), closest.height());
                vigra::MultiArray<2, f32_t> magnitudes(size);
                vigra::MultiArray<2, f32_t> orientations(size);
                for (u16_t x = 0; x < 1; x++) {
                    for (u16_t y = 0; y < closest.height(); y++) {
                        Point p(x, y);
                        magnitudes(x, y) = alg::gradientMagnitude(closest, p);
                        orientations(x, y) = alg::gradientOrientation(closest, p);
                    }
                }
                for (u16_t x = 8; x < interestPoints(e, i).width() - 8; x++) {
                    for (u16_t y = 8; y < interestPoints(e, i).height() - 8; y++) {
                        if (interestPoints(e, i)(x, y) >= 0) {
                            auto topLeftCorner = vigra::Shape2(x - 8, y - 8);
                            auto bottomRightCorner = vigra::Shape2(x + 8, y + 8);

                            auto orientation_region = orientations.subarray(topLeftCorner, bottomRightCorner);
                            auto mag_region = magnitudes.subarray(topLeftCorner, bottomRightCorner);

                            auto histogram = alg::orientationHistogram(orientation_region, mag_region, scale);

                            _createPeak(histogram);

                            //TODO:: Parabola to the 3 values closest to the peak

                        }
                    }
                }
            }
        }
    }

    void Sift::_createPeak(std::array<f32_t, 36>& histo) {
        const f32_t max = *(std::max_element(histo.begin(), histo.end()));
        //allowed range(80% of max)
        const f32_t range = max / 5 * 4;
        std::for_each(histo.begin(), histo.end(), [&](f32_t& elem) { if (elem < range) elem = -1; });
    }

    f32_t Sift::_calculateScale(u16_t e, u16_t i) const {
        if (e < 1) {
            if (i < 1)  {
                return _sigma;
            }
            return std::pow(_k, i) * _sigma;
        } else if (e == 1) {
            return std::pow(_k, _dogsPerEpoch + i) * _sigma;
        }
        //_dogsPerEpoch + 2: count of gaussians per epoch
        //e - 1: the very last epoch before our
        // - 2: take the top img from the stack and add i to it
        return std::pow(_k, (_dogsPerEpoch + 2) * (e - 1) - 2 + i) * _sigma;
    }

    void Sift::_eliminateEdgeResponses(Matrix<Matrix<f32_t>>& interestPoints,
            const Matrix<vigra::MultiArray<2, f32_t>>& dogs) const {

        for(u16_t e = 0; e < dogs.width(); e++) {
            for (u16_t i = 1; i < dogs.height() - 1; i++) {
                for (u16_t x = 1; x < dogs(e, i).shape(0) - 1; x++) {
                    for (u16_t y = 1; y < dogs(e, i).shape(1) - 1; y++) {
                        if (interestPoints(e, i - 1)(x, y) > -1) {
                            auto d = dogs(e, i);

                            const vigra::MultiArray<2, f32_t> param[3] = {dogs(e, i - 1), dogs(e, i), dogs(e, i + 1)};
                            const Point p(x, y);
                            vigra::Matrix<f32_t> deriv = alg::foDerivative(param, p);
                            vigra::Matrix<f32_t> sec_deriv = alg::soDerivative(param, p);

                            vigra::Matrix<f32_t> neg_sec_deriv = sec_deriv ;
                            neg_sec_deriv *=  -1;

                            vigra::MultiArray<2, f32_t> extremum(vigra::Shape2(3, 1));
                            if (!linearSolve(inverse(neg_sec_deriv), deriv, extremum)) {
                                std::cerr << "Couldn't solve linear system" << std::endl;
                                throw;
                            }

                            //Calculated up 0.5 from paper to own image values [0,255]
                            if (extremum(0, 0) > 127.5 || extremum(1, 0) > 127.5 || extremum(2, 0) > 127.5) {
                                interestPoints(e, i - 1)(x, y) = -1;
                                continue;
                            }
                            vigra::Matrix<f32_t> deriv_transpose = deriv.transpose();
                            f32_t func_val_extremum = dot(deriv_transpose, extremum);
                            func_val_extremum *= 0.5 + d(x, y);

                            //Calculated up 0.03 from paper to own image values[0, 255]
                            if (func_val_extremum < 7.65) {
                                interestPoints(e, i - 1)(x, y) = -1;
                                continue;
                            }

                            //dxx + dyy
                            f32_t hessian_tr = sec_deriv(0, 0) + sec_deriv(1, 1);
                            //dxx * dyy - dxy^2
                            f32_t hessian_det = sec_deriv(0, 0) * sec_deriv(1, 1) - std::pow(sec_deriv(0, 1), 2);

                            if (hessian_det < 0) {
                                interestPoints(e, i - 1)(x, y) = -1;
                                continue;
                            }

                            //Original r = 10, calculated up to own image values[0, 255]
                            //Question: The value 10 isn't based on the greyvalues?
                            if (std::pow(hessian_tr, 2) / hessian_det > std::pow(10 + 1, 2) / 10) {
                                interestPoints(e, i - 1)(x, y) = -1;
                            }
                        }
                    }
                }
            }
        }
    }

    const Matrix<Matrix<f32_t>> Sift::_findScaleSpaceExtrema(const Matrix<vigra::MultiArray<2, f32_t>>& dogs) const {
        //A matrix of matrix. Outer dogs will be ignored, because we need a upper and lower neighbor
        Matrix<Matrix<f32_t>> interestPoints(dogs.width(), dogs.height() - 2);
        for (u16_t e = 0; e < dogs.width(); e++) {
            for (u16_t i = 1; i < dogs.height() - 1; i++) {
                interestPoints(e, i - 1) = Matrix<f32_t>(dogs(e, i).width(), dogs(e, i).height(), -1);
                for (i16_t x = 1; x < dogs(e, i).width() - 1; x++) {
                    for (i16_t y = 1; y < dogs(e, i).height() - 1; y++) {
                        auto leftUpCorner = vigra::Shape2(x - 1, y - 1);
                        auto rightDownCorner = vigra::Shape2(x + 1, y + 1);

                        //Get the neighborhood of the current pixel
                        auto current = dogs(e, i).subarray(leftUpCorner, rightDownCorner);
                        //Get neighborhood of adjacent DOGs
                        auto under = dogs(e, i - 1).subarray(leftUpCorner, rightDownCorner);
                        auto above = dogs(e, i + 1).subarray(leftUpCorner, rightDownCorner);
                        //Check all neighborhood pixels of current and adjacent DOGs. If there isn't any
                        //pixel bigger or smaller than the current. We found an extremum.
                        if ((!any(current > dogs(e, i)(x, y)) &&
                                    !any(under > dogs(e, i)(x, y)) &&
                                    !any(above > dogs(e, i)(x, y))) ||
                                (!any(current < dogs(e, i)(x, y)) &&
                                 !any(under < dogs(e, i)(x, y)) &&
                                 !any(above < dogs(e, i)(x, y))))
                        {
                            interestPoints(e, i - 1)(x, y) = dogs(e, i)(x, y);
                        }
                    }
                }
            }
        }
        return interestPoints; // TODO: by ref entgegen nehmen, um copy zu vermeiden?
    }

    const Matrix<vigra::MultiArray<2, f32_t>> Sift::_createDOGs(vigra::MultiArray<2, f32_t>& img) {

        assert(_epochs > 0); // pre condition
        assert(_dogsPerEpoch >= 3); // pre condition

        Matrix<vigra::MultiArray<2, f32_t>> gaussians(_epochs, _dogsPerEpoch + 2);
        Matrix<vigra::MultiArray<2, f32_t>> dogs(_epochs, _dogsPerEpoch);

        gaussians(0, 0) = alg::convolveWithGauss(img, _sigma);

        //TODO: More elegant way?
        u16_t exp = 0;
        for (i16_t i = 0; i < _epochs; i++) {
            for (i16_t j = 1; j < _dogsPerEpoch + 1; j++) {
                gaussians(i, j) = alg::convolveWithGauss(gaussians(i, j - 1), std::pow(_k, exp) * _sigma);
                dogs(i, j - 1) = alg::dog(gaussians(i, j - 1), gaussians(i, j));
                exp++;
            }
            /*
             * If we aren't in the last epoch populate the next level with the second
             * last element, scaled by a half, of images of current epoch.
             */
            if (i < (_epochs - 1)) {
                auto scaledElem = alg::reduceToNextLevel(gaussians(i, _dogsPerEpoch - 1), _sigma);
                gaussians(i + 1, 0) = scaledElem;

                exp -= 2;
            }
        }
        _gaussians = gaussians;
        return dogs; // TODO: by ref entgegen nehmen um copy zu vermeiden?
    }
}
<|endoftext|>


<commit_before>#define _CRT_SECURE_NO_WARNINGS
#include <unordered_map>
#include <string>
#include <sstream>
#include <iomanip>
#include "imgui.h"
#define STB_LOG_IMPLEMENTATION
#include "stb_log.h"
#include "util.h"
#include "shellcmd.h"

class CImLog : public CCustomLog
{
public:
	CImLog() : m_cur_buf(nullptr) {
	}

	virtual ~CImLog() {
		while(m_cur_buf) {
			BufferHead *head = m_cur_buf;
			m_cur_buf = m_cur_buf->_next;
			delete [] head;
		}
	}

	typedef std::vector<std::pair<const char*, size_t>>::const_iterator iterator;

	inline iterator begin() const {
		return m_log_data.begin();
	}

	inline iterator end() const {
		return m_log_data.end();
	}

protected:
	virtual std::pair<char*, size_t> getstr(size_t required_size) override
	{
		if(!m_cur_buf || m_cur_buf->buf + required_size > m_cur_buf->end)
		{
			size_t size = required_size + sizeof(BufferHead);
			size = wyc::next_power2(size);
			char *new_buf = new char[size];
			BufferHead *head = (BufferHead*)new_buf;
			head->buf = new_buf;
			head->end = new_buf + size;
			if(m_cur_buf)
				head->_next = m_cur_buf;
			else
				head->_next = nullptr;
			m_cur_buf = head;
		}
		return {m_cur_buf->buf, m_cur_buf->end - m_cur_buf->buf};
	}

	virtual void setstr(size_t size) override {
		if(!m_cur_buf || m_cur_buf->buf + size > m_cur_buf->end) {
			assert(0);
			return;
		}
		size_t len = size - 1;
		if(m_cur_buf->buf[len] != 0) {
			assert(0);
			m_cur_buf->buf[len] = 0;
		}
		m_log_data.emplace_back(m_cur_buf->buf, len);
		m_cur_buf->buf += size;
	}

private:
	struct BufferHead {
		char *buf;
		char *end;
		BufferHead *_next;
	};
	BufferHead *m_cur_buf;
	std::vector<std::pair<const char*, size_t>> m_log_data;
};

class CImConsole
{
public:
	static inline CImConsole& singleton()
	{
		static CImConsole s_console;
		return s_console;
	}

	bool draw()
	{
		ImGui::SetNextWindowSize(ImVec2(float(m_width), float(m_height)), ImGuiCond_Always);
		if (!ImGui::Begin("console", 0, m_flags))
		{
			ImGui::End();
			return false;
		}

		const auto &style = ImGui::GetStyle();
		float progress_bar_height = 4;
		ImGui::SetCursorPosX(style.WindowPadding.x * 0.5f);
		ImGui::ProgressBar(0, ImVec2(ImGui::GetWindowWidth() - style.WindowPadding.x, progress_bar_height), "");
		ImGui::Separator();

		// BEGIN log
		ImGui::BeginChild("console_log", ImVec2(0, -ImGui::GetItemsLineHeightWithSpacing() - progress_bar_height - style.ItemSpacing.y),
			false, ImGuiWindowFlags_HorizontalScrollbar);
		if(m_log) {
			m_log->process();
			for(auto s : *m_log) {
				draw_log(s.first);
			}
		}
		ImGui::EndChild(); // END log

		// BEGIN input
		ImGui::Separator();
		ImGui::PushItemWidth(ImGui::GetContentRegionAvailWidth());
		if (ImGui::InputText("", m_input_beg, m_input_max, ImGuiInputTextFlags_EnterReturnsTrue | ImGuiInputTextFlags_CallbackCompletion | ImGuiInputTextFlags_CallbackHistory,
			[] (ImGuiTextEditCallbackData *ctx) -> int {
			((CImConsole*)ctx)->on_input_end();
			return 0;
		}, (void*)this))
		{
			if (m_input_beg[0]) {
				log_info(std::string(m_input_buf));
				process_input();
			}
			m_input_beg[0] = 0;
		}
		ImGui::PopItemWidth();

		 // Demonstrate keeping auto focus on the input box
		if (ImGui::IsItemHovered() || (ImGui::IsRootWindowOrAnyChildFocused() && !ImGui::IsAnyItemActive() && !ImGui::IsMouseClicked(0)))
			ImGui::SetKeyboardFocusHere(-1); // Auto focus previous widget
		// END input

		ImGui::End(); // END console
		return true;
	}

	void on_input_end()
	{
	}

	void process_input()
	{
		std::string cmd_name = strtok(m_input_beg, " ");
		if (cmd_name == "help")
		{
			int max_len = 4;
			for (auto &it : m_commands)
			{
				int s = it.first.size();
				if (s > max_len)
					max_len = s;
			}
			// tab size 4
			int c = (max_len + 4) & ~3;
			std::stringstream ss;
			ss << std::left << std::setfill(' ');
			ss << "- " << std::setw(c) << "help" << "show help message" << std::setw(0);
			log_info(ss.str().c_str());
			ss.str("");
			for (auto &it : m_commands) {
				ss << "- " << std::setw(c) << it.first << it.second->description() << std::setw(0);
				log_info(ss.str().c_str());
				ss.str("");
			}
			return;
		}
		auto iter = m_commands.find(cmd_name);
		if (iter == m_commands.end()) {
			log_error("Unkonwn command");
			return;
		}
		iter->second->execute(m_input_beg + cmd_name.size() + 1);
	}

	void draw_log(const std::string &str)
	{
		static const std::pair<std::string, ImVec4> s_tag_color[] = {
			std::make_pair(std::string("# "), ImVec4(0.0f, 1.0f, 0.0f, 1.0f)),
			std::make_pair(std::string("[ERROR]") , ImVec4(1.0f, 0.4f, 0.4f, 1.0f)),
			std::make_pair(std::string("[WARNING]"), ImVec4(1.0f, 1.0f, 0.0f, 1.0f)),
		};
		for (auto &tag : s_tag_color) {
			if (!str.compare(0, tag.first.size(), tag.first)) {
				ImGui::PushStyleColor(ImGuiCol_Text, tag.second);
				ImGui::TextUnformatted(str.c_str());
				ImGui::PopStyleColor();
				return;
			}
		}
		ImGui::TextUnformatted(str.c_str());
	}

	bool add_command(wyc::IShellCommand *cmd)
	{
		if (!cmd)
			return false;
		auto ret = m_commands.emplace(std::make_pair(cmd->name(), cmd));
		if (!ret.second)
			return false;
		return true;
	}

	inline void del_command(const std::string &cmd_name)
	{
		m_commands.erase(cmd_name);
	}

private:
	int m_width, m_height;
	ImGuiWindowFlags m_flags;
	static constexpr size_t INPUT_BUFF_SIZE = 256;
	char m_input_buf[INPUT_BUFF_SIZE];
	char *m_input_beg;
	size_t m_input_max;
	std::unordered_map<std::string, wyc::IShellCommand*> m_commands;
	CImLog *m_log;

	CImConsole()
	{
//		m_width = int(AppConfig::window_width * 0.382f), m_height = int(AppConfig::window_height);
		m_width = 200;
		m_height = 400;
		m_flags = ImGuiWindowFlags_NoResize
			| ImGuiWindowFlags_NoMove;
		const char *prompt = "# ";
		size_t prompt_len = strlen(prompt);
		strncpy(m_input_buf, prompt, INPUT_BUFF_SIZE);
		m_input_beg = m_input_buf + prompt_len;
		m_input_max = INPUT_BUFF_SIZE - prompt_len;
		m_log = new CImLog;
		add_log_handler(m_log);
	}

	~CImConsole()
	{
		if(m_log)
			delete m_log;
	}
};

bool console_register_command(wyc::IShellCommand *cmd)
{
	auto &console = CImConsole::singleton();
	return console.add_command(cmd);
}

void console_unregister_command(const std::string &cmd_name)
{
	auto &console = CImConsole::singleton();
	console.del_command(cmd_name);
}

bool show_console()
{
	auto &console = CImConsole::singleton();
	return console.draw();
}

<commit_msg>fix log buffer error<commit_after>#define _CRT_SECURE_NO_WARNINGS
#include <unordered_map>
#include <string>
#include <sstream>
#include <iomanip>
#include "imgui.h"
#define STB_LOG_IMPLEMENTATION
#include "stb_log.h"
#include "util.h"
#include "shellcmd.h"

class CImLog : public CCustomLog
{
public:
	CImLog() : m_cur_buf(nullptr) {
	}

	virtual ~CImLog() {
		while(m_cur_buf) {
			BufferHead *head = m_cur_buf;
			m_cur_buf = m_cur_buf->_next;
			delete [] head;
		}
	}

	typedef std::vector<std::pair<const char*, size_t>>::const_iterator iterator;

	inline iterator begin() const {
		return m_log_data.begin();
	}

	inline iterator end() const {
		return m_log_data.end();
	}

protected:
	virtual std::pair<char*, size_t> getstr(size_t required_size) override
	{
		if(!m_cur_buf || m_cur_buf->buf + required_size > m_cur_buf->end)
		{
			size_t size = required_size + sizeof(BufferHead);
			size = wyc::next_power2(size);
			char *new_buf = new char[size];
			BufferHead *head = (BufferHead*)new_buf;
			head->buf = new_buf + sizeof(BufferHead);
			head->end = new_buf + size;
			if(m_cur_buf)
				head->_next = m_cur_buf;
			else
				head->_next = nullptr;

			m_cur_buf = head;
		}
		return {m_cur_buf->buf, m_cur_buf->end - m_cur_buf->buf};
	}

	virtual void setstr(size_t size) override {
		if(!m_cur_buf || m_cur_buf->buf + size > m_cur_buf->end) {
			assert(0);
			return;
		}
		size_t len = size - 1;
		if(m_cur_buf->buf[len] != 0) {
			assert(0);
			m_cur_buf->buf[len] = 0;
		}
		m_log_data.emplace_back(m_cur_buf->buf, len);
		m_cur_buf->buf += size;
	}

private:
	struct BufferHead {
		char *buf;
		char *end;
		BufferHead *_next;
	};
	BufferHead *m_cur_buf;
	std::vector<std::pair<const char*, size_t>> m_log_data;
};

class CImConsole
{
public:
	static inline CImConsole& singleton()
	{
		static CImConsole s_console;
		return s_console;
	}

	bool draw()
	{
		ImGui::SetNextWindowSize(ImVec2(float(m_width), float(m_height)), ImGuiCond_Always);
		if (!ImGui::Begin("console", 0, m_flags))
		{
			ImGui::End();
			return false;
		}

		const auto &style = ImGui::GetStyle();
		float progress_bar_height = 4;
		ImGui::SetCursorPosX(style.WindowPadding.x * 0.5f);
		ImGui::ProgressBar(0, ImVec2(ImGui::GetWindowWidth() - style.WindowPadding.x, progress_bar_height), "");
		ImGui::Separator();

		// BEGIN log
		ImGui::BeginChild("console_log", ImVec2(0, -ImGui::GetItemsLineHeightWithSpacing() - progress_bar_height - style.ItemSpacing.y),
			false, ImGuiWindowFlags_HorizontalScrollbar);
		if(m_log) {
			m_log->process();
			for(auto s : *m_log) {
				draw_log(s.first);
			}
		}
		ImGui::EndChild(); // END log

		// BEGIN input
		ImGui::Separator();
		ImGui::PushItemWidth(ImGui::GetContentRegionAvailWidth());
		if (ImGui::InputText("", m_input_beg, m_input_max, ImGuiInputTextFlags_EnterReturnsTrue | ImGuiInputTextFlags_CallbackCompletion | ImGuiInputTextFlags_CallbackHistory,
			[] (ImGuiTextEditCallbackData *ctx) -> int {
			((CImConsole*)ctx)->on_input_end();
			return 0;
		}, (void*)this))
		{
			if (m_input_beg[0]) {
				log_info("%s", std::string(m_input_buf));
				process_input();
			}
			m_input_beg[0] = 0;
		}
		ImGui::PopItemWidth();

		 // Demonstrate keeping auto focus on the input box
		if (ImGui::IsItemHovered() || (ImGui::IsRootWindowOrAnyChildFocused() && !ImGui::IsAnyItemActive() && !ImGui::IsMouseClicked(0)))
			ImGui::SetKeyboardFocusHere(-1); // Auto focus previous widget
		// END input

		ImGui::End(); // END console
		return true;
	}

	void on_input_end()
	{
	}

	void process_input()
	{
		std::string cmd_name = strtok(m_input_beg, " ");
		if (cmd_name == "help")
		{
			int max_len = 4;
			for (auto &it : m_commands)
			{
				int s = it.first.size();
				if (s > max_len)
					max_len = s;
			}
			// tab size 4
			int c = (max_len + 4) & ~3;
			std::stringstream ss;
			ss << std::left << std::setfill(' ');
			ss << "- " << std::setw(c) << "help" << "show help message" << std::setw(0);
			log_info(ss.str().c_str());
			ss.str("");
			for (auto &it : m_commands) {
				ss << "- " << std::setw(c) << it.first << it.second->description() << std::setw(0);
				log_info(ss.str().c_str());
				ss.str("");
			}
			return;
		}
		auto iter = m_commands.find(cmd_name);
		if (iter == m_commands.end()) {
			log_error("Unkonwn command");
			return;
		}
		iter->second->execute(m_input_beg + cmd_name.size() + 1);
	}

	void draw_log(const std::string &str)
	{
		static const std::pair<std::string, ImVec4> s_tag_color[] = {
			std::make_pair(std::string("# "), ImVec4(0.0f, 1.0f, 0.0f, 1.0f)),
			std::make_pair(std::string("[ERROR]") , ImVec4(1.0f, 0.4f, 0.4f, 1.0f)),
			std::make_pair(std::string("[WARNING]"), ImVec4(1.0f, 1.0f, 0.0f, 1.0f)),
		};
		for (auto &tag : s_tag_color) {
			if (!str.compare(0, tag.first.size(), tag.first)) {
				ImGui::PushStyleColor(ImGuiCol_Text, tag.second);
				ImGui::TextUnformatted(str.c_str());
				ImGui::PopStyleColor();
				return;
			}
		}
		ImGui::TextUnformatted(str.c_str());
	}

	bool add_command(wyc::IShellCommand *cmd)
	{
		if (!cmd)
			return false;
		auto ret = m_commands.emplace(std::make_pair(cmd->name(), cmd));
		if (!ret.second)
			return false;
		return true;
	}

	inline void del_command(const std::string &cmd_name)
	{
		m_commands.erase(cmd_name);
	}

private:
	int m_width, m_height;
	ImGuiWindowFlags m_flags;
	static constexpr size_t INPUT_BUFF_SIZE = 256;
	char m_input_buf[INPUT_BUFF_SIZE];
	char *m_input_beg;
	size_t m_input_max;
	std::unordered_map<std::string, wyc::IShellCommand*> m_commands;
	CImLog *m_log;

	CImConsole()
	{
//		m_width = int(AppConfig::window_width * 0.382f), m_height = int(AppConfig::window_height);
		m_width = 200;
		m_height = 400;
		m_flags = ImGuiWindowFlags_NoResize
			| ImGuiWindowFlags_NoMove;
		const char *prompt = "# ";
		size_t prompt_len = strlen(prompt);
		strncpy(m_input_buf, prompt, INPUT_BUFF_SIZE);
		m_input_beg = m_input_buf + prompt_len;
		m_input_max = INPUT_BUFF_SIZE - prompt_len;
		m_log = new CImLog;
		add_log_handler(m_log);
	}

	~CImConsole()
	{
		if(m_log)
			delete m_log;
	}
};

bool console_register_command(wyc::IShellCommand *cmd)
{
	auto &console = CImConsole::singleton();
	return console.add_command(cmd);
}

void console_unregister_command(const std::string &cmd_name)
{
	auto &console = CImConsole::singleton();
	console.del_command(cmd_name);
}

bool show_console()
{
	auto &console = CImConsole::singleton();
	return console.draw();
}

<|endoftext|>


<commit_before>
#include <cstdlib>
#include <iostream>

#include "SupervisedLDA.hpp"
#include "ProgressVisitor.hpp"

int main(int argc, char **argv) {

    if (argc < 3) {
        std::cerr << "Usage:\n\t" << argv[0] << " TOPICS ITER" << std::endl;
        return 1;
    }

    SupervisedLDA<double> lda(std::atoi(argv[1]), std::atoi(argv[2]));

    int D, V;
    std::cin >> D >> V;

    MatrixXi X(V, D);
    VectorXi y(D);

    for (int d=0; d<D; d++) {
        std::cin >> y(d);
        for (int v=0; v<V; v++) {
            std::cin >> X(v, d);
        }
    }

    double likelihood;
    bool expectation;
    lda.set_progress_visitor(std::make_shared<FunctionVisitor<double> >([&](Progress<double> p){
        switch (p.state) {
            case Expectation:
                expectation = true;
                likelihood = p.value;
                if ((p.partial_iteration+1) % 100 == 0) {
                    std::cout << p.partial_iteration+1 << std::endl;
                }
                break;
            case Maximization:
                if (expectation) {
                    std::cout << "Likelihood: " << likelihood << std::endl;
                }
                expectation = false;
                std::cout << "log p(y | \\bar{z}, eta): " << -p.value << std::endl;
                break;
        }
    }));

    lda.fit(X, y);

    return 0;
}
<commit_msg>Add inference in the example binary<commit_after>
#include <cstdlib>
#include <iostream>

#include "SupervisedLDA.hpp"
#include "ProgressVisitor.hpp"

void split_training_test_set(
    const MatrixXi &X,
    const VectorXi &y,
    MatrixXi &train_X,
    MatrixXi &test_X,
    VectorXi &train_Y,
    VectorXi &test_Y)
{
    auto n = X.cols()/2;
    train_X.array() = X.block(0, 0, X.rows(), n);
    test_X.array() = X.block(0, n, X.rows(), X.cols()-n);
    train_Y.array() = y.segment(0, n);
    train_Y.array() = y.segment(n, n);

}

double accuracy_score(const VectorXi &y_true, const VectorXi &y_pred) {
    double accuracy = 0.0;

    for (int i=0; i<y_pred.rows(); i++) {
        if (y_pred(i) == y_true(i)) {
            accuracy += 1.0/y_pred.rows();
        }
    }

    return accuracy;
}

int main(int argc, char **argv) {

    if (argc < 3) {
        std::cerr << "Usage:\n\t" << argv[0] << " TOPICS ITER" << std::endl;
        return 1;
    }

    SupervisedLDA<double> lda(std::atoi(argv[1]), std::atoi(argv[2]));

    int D, V;
    std::cin >> D >> V;

    MatrixXi X(V, D);
    VectorXi y(D);

    for (int d=0; d<D; d++) {
        std::cin >> y(d);
        for (int v=0; v<V; v++) {
            std::cin >> X(v, d);
        }
    }
    MatrixXi X_train(V, D/2);
    VectorXi y_train(D/2);
    MatrixXi X_test(V, D/2);
    VectorXi y_test(D/2);

    split_training_test_set(X, y, X_train, X_test, y_train, y_test);

    double likelihood;
    bool expectation;
    lda.set_progress_visitor(std::make_shared<FunctionVisitor<double> >([&](Progress<double> p){
        switch (p.state) {
            case Expectation:
                expectation = true;
                likelihood = p.value;
                if ((p.partial_iteration+1) % 100 == 0) {
                    std::cout << p.partial_iteration+1 << std::endl;
                }
                break;
            case Maximization:
                if (expectation) {
                    std::cout << "Likelihood: " << likelihood << std::endl;
                }
                expectation = false;
                std::cout << "log p(y | \\bar{z}, eta): " << -p.value << std::endl;
                break;
        }
    }));

    lda.fit(X_train, y_train);
    VectorXi train_predictions = lda.predict(X_train);
    VectorXi test_predictions = lda.predict(X_test);

    std::cout << "Training Accuracy: " << accuracy_score(y_train, train_predictions) << std::endl;
    std::cout << "Test Accuracy: " << accuracy_score(y_test, test_predictions) << std::endl;

    return 0;
}
<|endoftext|>


<commit_before>// Copyright 2019 The TCMalloc Authors
//
// 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
//
//     https://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 "tcmalloc/span.h"

#include <stdint.h>

#include <algorithm>

#include "absl/base/optimization.h"
#include "absl/numeric/bits.h"
#include "tcmalloc/common.h"
#include "tcmalloc/internal/atomic_stats_counter.h"
#include "tcmalloc/internal/logging.h"
#include "tcmalloc/page_heap_allocator.h"
#include "tcmalloc/pagemap.h"
#include "tcmalloc/pages.h"
#include "tcmalloc/sampler.h"
#include "tcmalloc/static_vars.h"

GOOGLE_MALLOC_SECTION_BEGIN
namespace tcmalloc {
namespace tcmalloc_internal {

void Span::Sample(SampledAllocation* sampled_allocation) {
  ASSERT(!sampled_ && sampled_allocation);
  sampled_ = 1;
  sampled_allocation_ = sampled_allocation;

  // The cast to value matches Unsample.
  tcmalloc_internal::StatsCounter::Value allocated_bytes =
      static_cast<tcmalloc_internal::StatsCounter::Value>(
          AllocatedBytes(sampled_allocation->sampled_stack, true));
  Static::sampled_objects_size_.Add(allocated_bytes);
}

SampledAllocation* Span::Unsample() {
  if (!sampled_) {
    return nullptr;
  }
  sampled_ = 0;
  SampledAllocation* sampled_allocation = sampled_allocation_;
  sampled_allocation_ = nullptr;

  // The cast to Value ensures no funny business happens during the negation if
  // sizeof(size_t) != sizeof(Value).
  tcmalloc_internal::StatsCounter::Value neg_allocated_bytes =
      -static_cast<tcmalloc_internal::StatsCounter::Value>(
          AllocatedBytes(sampled_allocation->sampled_stack, true));
  Static::sampled_objects_size_.Add(neg_allocated_bytes);
  return sampled_allocation;
}

double Span::Fragmentation() const {
  const size_t size_class = Static::pagemap().sizeclass(first_page_);
  if (size_class == 0) {
    // Avoid crashes in production mode code, but report in tests.
    ASSERT(size_class != 0);
    return 0;
  }
  const size_t obj_size = Static::sizemap().class_to_size(size_class);
  const size_t span_objects = bytes_in_span() / obj_size;
  const size_t live = allocated_;
  if (live == 0) {
    // Avoid crashes in production mode code, but report in tests.
    ASSERT(live != 0);
    return 0;
  }
  // Assume that all in-use objects in this span are spread evenly
  // through this span.  So charge the free space in span evenly
  // to each of the live objects.
  // A note on units here: StackTraceTable::AddTrace(1, *t)
  // represents usage (of whatever kind: heap space, allocation,
  // fragmentation) of 1 object of size t->allocated_size.
  // So we want to report here the number of objects we are "responsible"
  // for pinning - NOT bytes.
  return static_cast<double>(span_objects - live) / live;
}

void Span::AverageFreelistAddedTime(const Span* other) {
  // Do this computation as floating-point to avoid overflowing our uint64_t.
  freelist_added_time_ = static_cast<uint64_t>(
      (static_cast<double>(freelist_added_time_) * num_pages_ +
       static_cast<double>(other->freelist_added_time_) * other->num_pages_) /
      (num_pages_ + other->num_pages_));
}

// Freelist organization.
//
// Partially full spans in CentralFreeList contain a list of free objects
// (freelist). We could use the free objects as linked list nodes and form
// a stack, but since the free objects are not likely to be cache-hot the
// chain of dependent misses is very cache-unfriendly. The current
// organization reduces number of cache misses during push/pop.
//
// Objects in the freelist are represented by 2-byte indices. The index is
// object offset from the span start divided by a constant. For small objects
// (<512) divider is 8, for larger -- 64. This allows to fit all indices into
// 2 bytes.
//
// The freelist has two components. First, we have a small array-based cache
// (4 objects) embedded directly into the Span (cache_ and cache_size_). We can
// access this without touching any objects themselves.
//
// The rest of the freelist is stored as arrays inside free objects themselves.
// We can store object_size / 2 indexes in any object, but this is not always
// sufficient to store the entire contents of a Span in a single object. So we
// reserve the first index slot in an object to form a linked list. We use the
// first object in that list (freelist_) as an array to push/pop from; any
// subsequent objects in the list's arrays are guaranteed to be full.
//
// Graphically this can be depicted as follows:
//
//         freelist_  embed_count_         cache_        cache_size_
// Span: [  |idx|         4          |idx|idx|---|---|        2      ]
//            |
//            \/
//            [idx|idx|idx|idx|idx|---|---|---]  16-byte object
//              |
//              \/
//              [---|idx|idx|idx|idx|idx|idx|idx]  16-byte object
//

Span::ObjIdx Span::PtrToIdx(void* ptr, size_t size) const {
  // Object index is an offset from span start divided by a power-of-two.
  // The divisors are choosen so that
  // (1) objects are aligned on the divisor,
  // (2) index fits into 16 bits and
  // (3) the index of the beginning of all objects is strictly less than
  //     kListEnd (note that we have 256K pages and multi-page spans).
  // For example with 1M spans we need kMultiPageAlignment >= 16.
  // An ASSERT in BuildFreelist() verifies a condition which implies (3).
  uintptr_t p = reinterpret_cast<uintptr_t>(ptr);
  uintptr_t off;
  if (size <= SizeMap::kMultiPageSize) {
    // Generally we need to load first_page_ to compute the offset.
    // But first_page_ can be in a different cache line then the fields that
    // we use in FreelistPush otherwise (cache_, cache_size_, freelist_).
    // So we avoid loading first_page_ for smaller sizes that have one page per
    // span, instead we compute the offset by taking low kPageShift bits of the
    // pointer.
    ASSERT(PageIdContaining(ptr) == first_page_);
    off = (p & (kPageSize - 1)) / kAlignment;
  } else {
    off = (p - first_page_.start_uintptr()) / SizeMap::kMultiPageAlignment;
  }
  ObjIdx idx = static_cast<ObjIdx>(off);
  ASSERT(idx != kListEnd);
  ASSERT(idx == off);
  return idx;
}

Span::ObjIdx* Span::IdxToPtr(ObjIdx idx, size_t size) const {
  ASSERT(idx != kListEnd);
  uintptr_t off = first_page_.start_uintptr() +
                  (static_cast<uintptr_t>(idx)
                   << (size <= SizeMap::kMultiPageSize
                           ? kAlignmentShift
                           : SizeMap::kMultiPageAlignmentShift));
  ObjIdx* ptr = reinterpret_cast<ObjIdx*>(off);
  ASSERT(PtrToIdx(ptr, size) == idx);
  return ptr;
}

Span::ObjIdx* Span::BitmapIdxToPtr(ObjIdx idx, size_t size) const {
  uintptr_t off =
      first_page_.start_uintptr() + (static_cast<uintptr_t>(idx) * size);
  ObjIdx* ptr = reinterpret_cast<ObjIdx*>(off);
  return ptr;
}

size_t Span::BitmapFreelistPopBatch(void** __restrict batch, size_t N,
                                    size_t size) {
#ifndef NDEBUG
  size_t before = bitmap_.CountBits(0, 64);
#endif  // NDEBUG

  size_t count = 0;
  // Want to fill the batch either with N objects, or the number of objects
  // remaining in the span.
  while (!bitmap_.IsZero() && count < N) {
    size_t offset = bitmap_.FindSet(0);
    ASSERT(offset < 64);
    batch[count] = BitmapIdxToPtr(offset, size);
    bitmap_.ClearLowestBit();
    count++;
  }

#ifndef NDEBUG
  size_t after = bitmap_.CountBits(0, 64);
  ASSERT(after + count == before);
  ASSERT(allocated_ + count == embed_count_ - after);
#endif  // NDEBUG
  allocated_ += count;
  return count;
}

size_t Span::FreelistPopBatch(void** __restrict batch, size_t N, size_t size) {
  // Handle spans with 64 or fewer objects using a bitmap. We expect spans
  // to frequently hold smaller objects.
  if (ABSL_PREDICT_FALSE(size >= kBitmapMinObjectSize)) {
    return BitmapFreelistPopBatch(batch, N, size);
  }
  if (ABSL_PREDICT_TRUE(size <= SizeMap::kMultiPageSize)) {
    return FreelistPopBatchSized<Align::SMALL>(batch, N, size);
  } else {
    return FreelistPopBatchSized<Align::LARGE>(batch, N, size);
  }
}

uint16_t Span::CalcReciprocal(size_t size) {
  // Calculate scaling factor. We want to avoid dividing by the size of the
  // object. Instead we'll multiply by a scaled version of the reciprocal.
  // We divide kBitmapScalingDenominator by the object size, so later we can
  // multiply by this reciprocal, and then divide this scaling factor out.
  // TODO(djgove) These divides can be computed once at start up.
  size_t reciprocal = 0;
  // The spans hold objects up to kMaxSize, so it's safe to assume.
  ABSL_ASSUME(size <= kMaxSize);
  if (size <= SizeMap::kMultiPageSize) {
    reciprocal = kBitmapScalingDenominator / (size >> kAlignmentShift);
  } else {
    reciprocal =
        kBitmapScalingDenominator / (size >> SizeMap::kMultiPageAlignmentShift);
  }
  ASSERT(reciprocal < 65536);
  return static_cast<uint16_t>(reciprocal);
}

void Span::BitmapBuildFreelist(size_t size, size_t count) {
  // We are using a bitmap to indicate whether objects are used or not. The
  // maximum capacity for the bitmap is 64 objects.
  ASSERT(count <= 64);
#ifndef NDEBUG
  // For bitmap_ use embed_count_ to record objects per span.
  embed_count_ = count;
#endif  // NDEBUG
  reciprocal_ = CalcReciprocal(size);
  allocated_ = 0;
  bitmap_.Clear();  // bitmap_ can be non-zero from a previous use.
  bitmap_.SetRange(0, count);
  ASSERT(bitmap_.CountBits(0, 64) == count);
}

int Span::BuildFreelist(size_t size, size_t count, void** batch, int N) {
  freelist_ = kListEnd;

  if (size >= kBitmapMinObjectSize) {
    BitmapBuildFreelist(size, count);
    return BitmapFreelistPopBatch(batch, N, size);
  }

  // First, push as much as we can into the batch.
  char* ptr = static_cast<char*>(start_address());
  int result = N <= count ? N : count;
  for (int i = 0; i < result; ++i) {
    batch[i] = ptr;
    ptr += size;
  }
  allocated_ = result;

  ObjIdx idxStep = size / kAlignment;
  // Valid objects are {0, idxStep, idxStep * 2, ..., idxStep * (count - 1)}.
  if (size > SizeMap::kMultiPageSize) {
    idxStep = size / SizeMap::kMultiPageAlignment;
  }
  ObjIdx idx = idxStep * result;

  // Verify that the end of the useful portion of the span (and the beginning of
  // the span waste) has an index that doesn't overflow or risk confusion with
  // kListEnd. This is slightly stronger than we actually need (see comment in
  // PtrToIdx for that) but rules out some bugs and weakening it wouldn't
  // actually help. One example of the potential bugs that are ruled out is the
  // possibility of idxEnd (below) overflowing.
  ASSERT(count * idxStep < kListEnd);

  // The index of the end of the useful portion of the span.
  ObjIdx idxEnd = count * idxStep;

  // Then, push as much as we can into the cache_.
  int cache_size = 0;
  for (; idx < idxEnd && cache_size < kCacheSize; idx += idxStep) {
    cache_[cache_size] = idx;
    cache_size++;
  }
  cache_size_ = cache_size;

  // Now, build freelist and stack other objects onto freelist objects.
  // Note: we take freelist objects from the beginning and stacked objects
  // from the end. This has a nice property of not paging in whole span at once
  // and not draining whole cache.
  ObjIdx* host = nullptr;  // cached first object on freelist
  const size_t max_embed = size / sizeof(ObjIdx) - 1;
  int embed_count = 0;
  while (idx < idxEnd) {
    // Check the no idx can be confused with kListEnd.
    ASSERT(idx != kListEnd);
    if (host && embed_count != max_embed) {
      // Push onto first object on the freelist.
      embed_count++;
      idxEnd -= idxStep;
      host[embed_count] = idxEnd;
    } else {
      // The first object is full, push new object onto freelist.
      host = IdxToPtr(idx, size);
      host[0] = freelist_;
      freelist_ = idx;
      embed_count = 0;
      idx += idxStep;
    }
  }
  embed_count_ = embed_count;
  return result;
}

}  // namespace tcmalloc_internal
}  // namespace tcmalloc
GOOGLE_MALLOC_SECTION_END
<commit_msg>Fail loudly when sampled_ and sampled_allocation_ is at conflicting states at the start of Span::UnSample().<commit_after>// Copyright 2019 The TCMalloc Authors
//
// 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
//
//     https://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 "tcmalloc/span.h"

#include <stdint.h>

#include <algorithm>

#include "absl/base/optimization.h"
#include "absl/numeric/bits.h"
#include "tcmalloc/common.h"
#include "tcmalloc/internal/atomic_stats_counter.h"
#include "tcmalloc/internal/logging.h"
#include "tcmalloc/page_heap_allocator.h"
#include "tcmalloc/pagemap.h"
#include "tcmalloc/pages.h"
#include "tcmalloc/sampler.h"
#include "tcmalloc/static_vars.h"

GOOGLE_MALLOC_SECTION_BEGIN
namespace tcmalloc {
namespace tcmalloc_internal {

void Span::Sample(SampledAllocation* sampled_allocation) {
  ASSERT(!sampled_ && sampled_allocation);
  sampled_ = 1;
  sampled_allocation_ = sampled_allocation;

  // The cast to value matches Unsample.
  tcmalloc_internal::StatsCounter::Value allocated_bytes =
      static_cast<tcmalloc_internal::StatsCounter::Value>(
          AllocatedBytes(sampled_allocation->sampled_stack, true));
  Static::sampled_objects_size_.Add(allocated_bytes);
}

SampledAllocation* Span::Unsample() {
  if (!sampled_) {
    return nullptr;
  }
  ASSERT(sampled_ && sampled_allocation_);
  sampled_ = 0;
  SampledAllocation* sampled_allocation = sampled_allocation_;
  sampled_allocation_ = nullptr;

  // The cast to Value ensures no funny business happens during the negation if
  // sizeof(size_t) != sizeof(Value).
  tcmalloc_internal::StatsCounter::Value neg_allocated_bytes =
      -static_cast<tcmalloc_internal::StatsCounter::Value>(
          AllocatedBytes(sampled_allocation->sampled_stack, true));
  Static::sampled_objects_size_.Add(neg_allocated_bytes);
  return sampled_allocation;
}

double Span::Fragmentation() const {
  const size_t size_class = Static::pagemap().sizeclass(first_page_);
  if (size_class == 0) {
    // Avoid crashes in production mode code, but report in tests.
    ASSERT(size_class != 0);
    return 0;
  }
  const size_t obj_size = Static::sizemap().class_to_size(size_class);
  const size_t span_objects = bytes_in_span() / obj_size;
  const size_t live = allocated_;
  if (live == 0) {
    // Avoid crashes in production mode code, but report in tests.
    ASSERT(live != 0);
    return 0;
  }
  // Assume that all in-use objects in this span are spread evenly
  // through this span.  So charge the free space in span evenly
  // to each of the live objects.
  // A note on units here: StackTraceTable::AddTrace(1, *t)
  // represents usage (of whatever kind: heap space, allocation,
  // fragmentation) of 1 object of size t->allocated_size.
  // So we want to report here the number of objects we are "responsible"
  // for pinning - NOT bytes.
  return static_cast<double>(span_objects - live) / live;
}

void Span::AverageFreelistAddedTime(const Span* other) {
  // Do this computation as floating-point to avoid overflowing our uint64_t.
  freelist_added_time_ = static_cast<uint64_t>(
      (static_cast<double>(freelist_added_time_) * num_pages_ +
       static_cast<double>(other->freelist_added_time_) * other->num_pages_) /
      (num_pages_ + other->num_pages_));
}

// Freelist organization.
//
// Partially full spans in CentralFreeList contain a list of free objects
// (freelist). We could use the free objects as linked list nodes and form
// a stack, but since the free objects are not likely to be cache-hot the
// chain of dependent misses is very cache-unfriendly. The current
// organization reduces number of cache misses during push/pop.
//
// Objects in the freelist are represented by 2-byte indices. The index is
// object offset from the span start divided by a constant. For small objects
// (<512) divider is 8, for larger -- 64. This allows to fit all indices into
// 2 bytes.
//
// The freelist has two components. First, we have a small array-based cache
// (4 objects) embedded directly into the Span (cache_ and cache_size_). We can
// access this without touching any objects themselves.
//
// The rest of the freelist is stored as arrays inside free objects themselves.
// We can store object_size / 2 indexes in any object, but this is not always
// sufficient to store the entire contents of a Span in a single object. So we
// reserve the first index slot in an object to form a linked list. We use the
// first object in that list (freelist_) as an array to push/pop from; any
// subsequent objects in the list's arrays are guaranteed to be full.
//
// Graphically this can be depicted as follows:
//
//         freelist_  embed_count_         cache_        cache_size_
// Span: [  |idx|         4          |idx|idx|---|---|        2      ]
//            |
//            \/
//            [idx|idx|idx|idx|idx|---|---|---]  16-byte object
//              |
//              \/
//              [---|idx|idx|idx|idx|idx|idx|idx]  16-byte object
//

Span::ObjIdx Span::PtrToIdx(void* ptr, size_t size) const {
  // Object index is an offset from span start divided by a power-of-two.
  // The divisors are choosen so that
  // (1) objects are aligned on the divisor,
  // (2) index fits into 16 bits and
  // (3) the index of the beginning of all objects is strictly less than
  //     kListEnd (note that we have 256K pages and multi-page spans).
  // For example with 1M spans we need kMultiPageAlignment >= 16.
  // An ASSERT in BuildFreelist() verifies a condition which implies (3).
  uintptr_t p = reinterpret_cast<uintptr_t>(ptr);
  uintptr_t off;
  if (size <= SizeMap::kMultiPageSize) {
    // Generally we need to load first_page_ to compute the offset.
    // But first_page_ can be in a different cache line then the fields that
    // we use in FreelistPush otherwise (cache_, cache_size_, freelist_).
    // So we avoid loading first_page_ for smaller sizes that have one page per
    // span, instead we compute the offset by taking low kPageShift bits of the
    // pointer.
    ASSERT(PageIdContaining(ptr) == first_page_);
    off = (p & (kPageSize - 1)) / kAlignment;
  } else {
    off = (p - first_page_.start_uintptr()) / SizeMap::kMultiPageAlignment;
  }
  ObjIdx idx = static_cast<ObjIdx>(off);
  ASSERT(idx != kListEnd);
  ASSERT(idx == off);
  return idx;
}

Span::ObjIdx* Span::IdxToPtr(ObjIdx idx, size_t size) const {
  ASSERT(idx != kListEnd);
  uintptr_t off = first_page_.start_uintptr() +
                  (static_cast<uintptr_t>(idx)
                   << (size <= SizeMap::kMultiPageSize
                           ? kAlignmentShift
                           : SizeMap::kMultiPageAlignmentShift));
  ObjIdx* ptr = reinterpret_cast<ObjIdx*>(off);
  ASSERT(PtrToIdx(ptr, size) == idx);
  return ptr;
}

Span::ObjIdx* Span::BitmapIdxToPtr(ObjIdx idx, size_t size) const {
  uintptr_t off =
      first_page_.start_uintptr() + (static_cast<uintptr_t>(idx) * size);
  ObjIdx* ptr = reinterpret_cast<ObjIdx*>(off);
  return ptr;
}

size_t Span::BitmapFreelistPopBatch(void** __restrict batch, size_t N,
                                    size_t size) {
#ifndef NDEBUG
  size_t before = bitmap_.CountBits(0, 64);
#endif  // NDEBUG

  size_t count = 0;
  // Want to fill the batch either with N objects, or the number of objects
  // remaining in the span.
  while (!bitmap_.IsZero() && count < N) {
    size_t offset = bitmap_.FindSet(0);
    ASSERT(offset < 64);
    batch[count] = BitmapIdxToPtr(offset, size);
    bitmap_.ClearLowestBit();
    count++;
  }

#ifndef NDEBUG
  size_t after = bitmap_.CountBits(0, 64);
  ASSERT(after + count == before);
  ASSERT(allocated_ + count == embed_count_ - after);
#endif  // NDEBUG
  allocated_ += count;
  return count;
}

size_t Span::FreelistPopBatch(void** __restrict batch, size_t N, size_t size) {
  // Handle spans with 64 or fewer objects using a bitmap. We expect spans
  // to frequently hold smaller objects.
  if (ABSL_PREDICT_FALSE(size >= kBitmapMinObjectSize)) {
    return BitmapFreelistPopBatch(batch, N, size);
  }
  if (ABSL_PREDICT_TRUE(size <= SizeMap::kMultiPageSize)) {
    return FreelistPopBatchSized<Align::SMALL>(batch, N, size);
  } else {
    return FreelistPopBatchSized<Align::LARGE>(batch, N, size);
  }
}

uint16_t Span::CalcReciprocal(size_t size) {
  // Calculate scaling factor. We want to avoid dividing by the size of the
  // object. Instead we'll multiply by a scaled version of the reciprocal.
  // We divide kBitmapScalingDenominator by the object size, so later we can
  // multiply by this reciprocal, and then divide this scaling factor out.
  // TODO(djgove) These divides can be computed once at start up.
  size_t reciprocal = 0;
  // The spans hold objects up to kMaxSize, so it's safe to assume.
  ABSL_ASSUME(size <= kMaxSize);
  if (size <= SizeMap::kMultiPageSize) {
    reciprocal = kBitmapScalingDenominator / (size >> kAlignmentShift);
  } else {
    reciprocal =
        kBitmapScalingDenominator / (size >> SizeMap::kMultiPageAlignmentShift);
  }
  ASSERT(reciprocal < 65536);
  return static_cast<uint16_t>(reciprocal);
}

void Span::BitmapBuildFreelist(size_t size, size_t count) {
  // We are using a bitmap to indicate whether objects are used or not. The
  // maximum capacity for the bitmap is 64 objects.
  ASSERT(count <= 64);
#ifndef NDEBUG
  // For bitmap_ use embed_count_ to record objects per span.
  embed_count_ = count;
#endif  // NDEBUG
  reciprocal_ = CalcReciprocal(size);
  allocated_ = 0;
  bitmap_.Clear();  // bitmap_ can be non-zero from a previous use.
  bitmap_.SetRange(0, count);
  ASSERT(bitmap_.CountBits(0, 64) == count);
}

int Span::BuildFreelist(size_t size, size_t count, void** batch, int N) {
  freelist_ = kListEnd;

  if (size >= kBitmapMinObjectSize) {
    BitmapBuildFreelist(size, count);
    return BitmapFreelistPopBatch(batch, N, size);
  }

  // First, push as much as we can into the batch.
  char* ptr = static_cast<char*>(start_address());
  int result = N <= count ? N : count;
  for (int i = 0; i < result; ++i) {
    batch[i] = ptr;
    ptr += size;
  }
  allocated_ = result;

  ObjIdx idxStep = size / kAlignment;
  // Valid objects are {0, idxStep, idxStep * 2, ..., idxStep * (count - 1)}.
  if (size > SizeMap::kMultiPageSize) {
    idxStep = size / SizeMap::kMultiPageAlignment;
  }
  ObjIdx idx = idxStep * result;

  // Verify that the end of the useful portion of the span (and the beginning of
  // the span waste) has an index that doesn't overflow or risk confusion with
  // kListEnd. This is slightly stronger than we actually need (see comment in
  // PtrToIdx for that) but rules out some bugs and weakening it wouldn't
  // actually help. One example of the potential bugs that are ruled out is the
  // possibility of idxEnd (below) overflowing.
  ASSERT(count * idxStep < kListEnd);

  // The index of the end of the useful portion of the span.
  ObjIdx idxEnd = count * idxStep;

  // Then, push as much as we can into the cache_.
  int cache_size = 0;
  for (; idx < idxEnd && cache_size < kCacheSize; idx += idxStep) {
    cache_[cache_size] = idx;
    cache_size++;
  }
  cache_size_ = cache_size;

  // Now, build freelist and stack other objects onto freelist objects.
  // Note: we take freelist objects from the beginning and stacked objects
  // from the end. This has a nice property of not paging in whole span at once
  // and not draining whole cache.
  ObjIdx* host = nullptr;  // cached first object on freelist
  const size_t max_embed = size / sizeof(ObjIdx) - 1;
  int embed_count = 0;
  while (idx < idxEnd) {
    // Check the no idx can be confused with kListEnd.
    ASSERT(idx != kListEnd);
    if (host && embed_count != max_embed) {
      // Push onto first object on the freelist.
      embed_count++;
      idxEnd -= idxStep;
      host[embed_count] = idxEnd;
    } else {
      // The first object is full, push new object onto freelist.
      host = IdxToPtr(idx, size);
      host[0] = freelist_;
      freelist_ = idx;
      embed_count = 0;
      idx += idxStep;
    }
  }
  embed_count_ = embed_count;
  return result;
}

}  // namespace tcmalloc_internal
}  // namespace tcmalloc
GOOGLE_MALLOC_SECTION_END
<|endoftext|>


<commit_before>/* -*- Mode: C++; tab-width: 8; indent-tabs-mode: t; c-basic-offset: 8 -*- */

/*
 *    Lorica source file.
 *    Copyright (C) 2007-2008 OMC Denmark ApS.
 *
 *    This program is free software; you can redistribute it and/or modify
 *    it under the terms of the GNU General Public License as published by
 *    the Free Software Foundation; either version 2 of the License, or
 *    (at your option) any later version.
 *
 *    This program is distributed in the hope that it will be useful,
 *    but WITHOUT ANY WARRANTY; without even the implied warranty of
 *    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *    GNU General Public License for more details.
 *
 *    You should have received a copy of the GNU General Public License
 *    along with this program; if not, write to the Free Software
 *    Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA 02111-1307 USA
 */

#include "defines/build-system-defs.h"

#include <ace/Dynamic_Service.h>
#include <tao/ORB_Core.h>

#include "lorica/MapperRegistry.h"
#include "lorica/ReferenceMapValue.h"
#include "lorica/debug.h"

#include "proxy/ReferenceMapper_i.h"

Lorica::ReferenceMapper_i::ReferenceMapper_i(CORBA::ORB_ptr orb,
					     IORTable::Table_ptr iort,
					     bool has_sec)
	: orb_(CORBA::ORB::_duplicate(orb)),
	  ior_table_(IORTable::Table::_duplicate(iort)),
	  has_security_(has_sec),
	  access_decision_(TAO::SL2::AccessDecision::_nil())
{
	this->registry_ = ACE_Dynamic_Service<Lorica_MapperRegistry>::instance("MapperRegistry");
	if (this->registry_ == 0 && Lorica_debug_level > 0)
		ACE_DEBUG((LM_DEBUG, ACE_TEXT("(%P|%t) %N:%l - no registry\n")));

	try {
		if (this->has_security_) {
			CORBA::Object_var obj =	orb_->resolve_initial_references("SecurityLevel2:SecurityManager");
			SecurityLevel2::SecurityManager_var sl2sm = SecurityLevel2::SecurityManager::_narrow(obj.in());
			SecurityLevel2::AccessDecision_var sl2ad = sl2sm->access_decision();
			this->access_decision_ = TAO::SL2::AccessDecision::_narrow(sl2ad.in());
		} else {
			if (Lorica_debug_level > 0)
				ACE_DEBUG((LM_DEBUG, ACE_TEXT("(%P|%t) %N:%l - SSLIOP not available\n")));
		}
	}
	catch (CORBA::Exception &) {
		ACE_ERROR ((LM_ERROR, "%N:%l - cannot initialize access decision object\n"));
	}
}

Lorica::ReferenceMapper_i::~ReferenceMapper_i(void)
{
}

void
Lorica::ReferenceMapper_i::allow_insecure_access(CORBA::Object_ptr self)
{
	if (this->has_security_ && !CORBA::is_nil(this->access_decision_))
	{
		CORBA::String_var orbid = this->orb_->id();
		PortableServer::ObjectId_var oid;
		CORBA::OctetSeq_var poaid;

		if (this->registry_->poa_and_object_id (self, poaid.out(), oid.out()))
			this->access_decision_->add_object(orbid.in(), poaid.in(), oid.in(),
							   true);
	}
}

CORBA::Object_ptr
Lorica::ReferenceMapper_i::as_secure_server(CORBA::Object_ptr orig,
					    const char *corbaloc_name,
					    Lorica::ServerAgent_ptr agent)
	throw (CORBA::SystemException,
	       Lorica::ReferenceMapper::SecurityNotAvailable)
{
	if (!this->has_security_)
		throw Lorica::ReferenceMapper::SecurityNotAvailable();

	if (CORBA::is_nil(orig))
		return CORBA::Object::_nil();

	return this->as_server_i(true, orig, corbaloc_name, agent);
}

CORBA::Object_ptr
Lorica::ReferenceMapper_i::as_server(CORBA::Object_ptr orig,
				     const char *corbaloc_name,
				     Lorica::ServerAgent_ptr agent)
	throw (CORBA::SystemException)
{
	if (CORBA::is_nil(orig))
		return CORBA::Object::_nil();

	return this->as_server_i (false, orig, corbaloc_name, agent);
}

CORBA::Object_ptr
Lorica::ReferenceMapper_i::as_server_i(bool require_secure,
				       CORBA::Object_ptr orig,
				       const char *corbaloc_name,
				       Lorica::ServerAgent_ptr agent)
	throw (CORBA::SystemException)
{
	Lorica::ReferenceMapValue_var rmv = this->registry_->map_reference(orig, true, require_secure);

	if (rmv.get() == 0) {
		if(Lorica_debug_level > 0)
			ACE_DEBUG((LM_DEBUG, ACE_TEXT("(%P|%t) %N:%l - registry returned a null ReferenceMapValue\n")));

		return CORBA::Object::_nil();
	}

	if (Lorica_debug_level > 5)
		ACE_DEBUG((LM_DEBUG, ACE_TEXT("(%P|%t) %N:%l - registry returned a non-null rmv\n")));

	if (corbaloc_name[0] != 0 && !CORBA::is_nil(this->ior_table_.in())) {
		CORBA::String_var mapped = this->orb_->object_to_string(rmv->mapped_ref_.in());
		if(Lorica_debug_level > 5) {
			ACE_DEBUG((LM_DEBUG,
				   ACE_TEXT("(%P|%t) %N:%l - adding mapped value to IORTable, %s, %s\n"),
				   corbaloc_name,mapped.in()));
		}

		try {
			ior_table_->bind (corbaloc_name,mapped.in());
			if (Lorica_debug_level > 5)
				ACE_DEBUG((LM_DEBUG, ACE_TEXT("(%P|%t) %N:%l - bound to ior table\n")));
			rmv->ior_table_name_ = corbaloc_name;
		}
		catch (IORTable::AlreadyBound &) {
			ACE_ERROR((LM_ERROR,
				   "(%P|%t) %N:%l - ior_table cannot bind to %s\n",
				   corbaloc_name));
		}
	}

	rmv->require_secure_ = require_secure;
	if (!require_secure)
		this->allow_insecure_access(rmv->mapped_ref_.in());

	if (!CORBA::is_nil(agent))
		rmv->agent_ = Lorica::ServerAgent::_duplicate(agent);

	return CORBA::Object::_duplicate(rmv->mapped_ref_.in());
}

CORBA::Object_ptr
Lorica::ReferenceMapper_i::as_client(CORBA::Object_ptr orig)
	throw (CORBA::SystemException)
{
	if (CORBA::is_nil(orig))
		return CORBA::Object::_nil();

	Lorica::ReferenceMapValue_var rmv = this->registry_->map_reference(orig, false);
	if (rmv.get() == 0)
		return CORBA::Object::_nil();

	return CORBA::Object::_duplicate(rmv->mapped_ref_);
}

void
Lorica::ReferenceMapper_i::remove_server(CORBA::Object_ptr mapped)
	throw (CORBA::SystemException)
{
	if (CORBA::is_nil(mapped))
		return CORBA::Object::_nil();

	Lorica::ReferenceMapValue_var rmv = this->registry_->remove_reference (mapped, true);
	if ((rmv.get() != 0) && (rmv->ior_table_name_ != "")) {
		if (Lorica_debug_level > 5) {
			ACE_DEBUG((LM_DEBUG,
				   ACE_TEXT("(%P|%t) %N:%l - unbinding %s from IOR Table\n"),
				   rmv->ior_table_name_.c_str()));
		}

		this->ior_table_->unbind(rmv->ior_table_name_.c_str());
	} else {
		if (Lorica_debug_level > 5)
			ACE_DEBUG((LM_DEBUG,
				   ACE_TEXT("(%P|%t) %N:%l - remove reference returned rmv = 0x%x\n"),
				   rmv.get()));
	}

	// rmv releases reference to value which cleans itself up.
}

void
Lorica::ReferenceMapper_i::remove_client(CORBA::Object_ptr mapped)
	throw (CORBA::SystemException)
{
	if (CORBA::is_nil(mapped))
		return CORBA::Object::_nil();

	Lorica::ReferenceMapValue_var rmv = this->registry_->remove_reference(mapped, false);
}

CORBA::Boolean
Lorica::ReferenceMapper_i::supports_security(void)
	throw (CORBA::SystemException)
{
	return this->has_security_;
}
<commit_msg>Fix build error<commit_after>/* -*- Mode: C++; tab-width: 8; indent-tabs-mode: t; c-basic-offset: 8 -*- */

/*
 *    Lorica source file.
 *    Copyright (C) 2007-2008 OMC Denmark ApS.
 *
 *    This program is free software; you can redistribute it and/or modify
 *    it under the terms of the GNU General Public License as published by
 *    the Free Software Foundation; either version 2 of the License, or
 *    (at your option) any later version.
 *
 *    This program is distributed in the hope that it will be useful,
 *    but WITHOUT ANY WARRANTY; without even the implied warranty of
 *    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *    GNU General Public License for more details.
 *
 *    You should have received a copy of the GNU General Public License
 *    along with this program; if not, write to the Free Software
 *    Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA 02111-1307 USA
 */

#include "defines/build-system-defs.h"

#include <ace/Dynamic_Service.h>
#include <tao/ORB_Core.h>

#include "lorica/MapperRegistry.h"
#include "lorica/ReferenceMapValue.h"
#include "lorica/debug.h"

#include "proxy/ReferenceMapper_i.h"

Lorica::ReferenceMapper_i::ReferenceMapper_i(CORBA::ORB_ptr orb,
					     IORTable::Table_ptr iort,
					     bool has_sec)
	: orb_(CORBA::ORB::_duplicate(orb)),
	  ior_table_(IORTable::Table::_duplicate(iort)),
	  has_security_(has_sec),
	  access_decision_(TAO::SL2::AccessDecision::_nil())
{
	this->registry_ = ACE_Dynamic_Service<Lorica_MapperRegistry>::instance("MapperRegistry");
	if (this->registry_ == 0 && Lorica_debug_level > 0)
		ACE_DEBUG((LM_DEBUG, ACE_TEXT("(%P|%t) %N:%l - no registry\n")));

	try {
		if (this->has_security_) {
			CORBA::Object_var obj =	orb_->resolve_initial_references("SecurityLevel2:SecurityManager");
			SecurityLevel2::SecurityManager_var sl2sm = SecurityLevel2::SecurityManager::_narrow(obj.in());
			SecurityLevel2::AccessDecision_var sl2ad = sl2sm->access_decision();
			this->access_decision_ = TAO::SL2::AccessDecision::_narrow(sl2ad.in());
		} else {
			if (Lorica_debug_level > 0)
				ACE_DEBUG((LM_DEBUG, ACE_TEXT("(%P|%t) %N:%l - SSLIOP not available\n")));
		}
	}
	catch (CORBA::Exception &) {
		ACE_ERROR ((LM_ERROR, "%N:%l - cannot initialize access decision object\n"));
	}
}

Lorica::ReferenceMapper_i::~ReferenceMapper_i(void)
{
}

void
Lorica::ReferenceMapper_i::allow_insecure_access(CORBA::Object_ptr self)
{
	if (this->has_security_ && !CORBA::is_nil(this->access_decision_))
	{
		CORBA::String_var orbid = this->orb_->id();
		PortableServer::ObjectId_var oid;
		CORBA::OctetSeq_var poaid;

		if (this->registry_->poa_and_object_id (self, poaid.out(), oid.out()))
			this->access_decision_->add_object(orbid.in(), poaid.in(), oid.in(),
							   true);
	}
}

CORBA::Object_ptr
Lorica::ReferenceMapper_i::as_secure_server(CORBA::Object_ptr orig,
					    const char *corbaloc_name,
					    Lorica::ServerAgent_ptr agent)
	throw (CORBA::SystemException,
	       Lorica::ReferenceMapper::SecurityNotAvailable)
{
	if (!this->has_security_)
		throw Lorica::ReferenceMapper::SecurityNotAvailable();

	if (CORBA::is_nil(orig))
		return CORBA::Object::_nil();

	return this->as_server_i(true, orig, corbaloc_name, agent);
}

CORBA::Object_ptr
Lorica::ReferenceMapper_i::as_server(CORBA::Object_ptr orig,
				     const char *corbaloc_name,
				     Lorica::ServerAgent_ptr agent)
	throw (CORBA::SystemException)
{
	if (CORBA::is_nil(orig))
		return CORBA::Object::_nil();

	return this->as_server_i (false, orig, corbaloc_name, agent);
}

CORBA::Object_ptr
Lorica::ReferenceMapper_i::as_server_i(bool require_secure,
				       CORBA::Object_ptr orig,
				       const char *corbaloc_name,
				       Lorica::ServerAgent_ptr agent)
	throw (CORBA::SystemException)
{
	Lorica::ReferenceMapValue_var rmv = this->registry_->map_reference(orig, true, require_secure);

	if (rmv.get() == 0) {
		if(Lorica_debug_level > 0)
			ACE_DEBUG((LM_DEBUG, ACE_TEXT("(%P|%t) %N:%l - registry returned a null ReferenceMapValue\n")));

		return CORBA::Object::_nil();
	}

	if (Lorica_debug_level > 5)
		ACE_DEBUG((LM_DEBUG, ACE_TEXT("(%P|%t) %N:%l - registry returned a non-null rmv\n")));

	if (corbaloc_name[0] != 0 && !CORBA::is_nil(this->ior_table_.in())) {
		CORBA::String_var mapped = this->orb_->object_to_string(rmv->mapped_ref_.in());
		if(Lorica_debug_level > 5) {
			ACE_DEBUG((LM_DEBUG,
				   ACE_TEXT("(%P|%t) %N:%l - adding mapped value to IORTable, %s, %s\n"),
				   corbaloc_name,mapped.in()));
		}

		try {
			ior_table_->bind (corbaloc_name,mapped.in());
			if (Lorica_debug_level > 5)
				ACE_DEBUG((LM_DEBUG, ACE_TEXT("(%P|%t) %N:%l - bound to ior table\n")));
			rmv->ior_table_name_ = corbaloc_name;
		}
		catch (IORTable::AlreadyBound &) {
			ACE_ERROR((LM_ERROR,
				   "(%P|%t) %N:%l - ior_table cannot bind to %s\n",
				   corbaloc_name));
		}
	}

	rmv->require_secure_ = require_secure;
	if (!require_secure)
		this->allow_insecure_access(rmv->mapped_ref_.in());

	if (!CORBA::is_nil(agent))
		rmv->agent_ = Lorica::ServerAgent::_duplicate(agent);

	return CORBA::Object::_duplicate(rmv->mapped_ref_.in());
}

CORBA::Object_ptr
Lorica::ReferenceMapper_i::as_client(CORBA::Object_ptr orig)
	throw (CORBA::SystemException)
{
	if (CORBA::is_nil(orig))
		return CORBA::Object::_nil();

	Lorica::ReferenceMapValue_var rmv = this->registry_->map_reference(orig, false);
	if (rmv.get() == 0)
		return CORBA::Object::_nil();

	return CORBA::Object::_duplicate(rmv->mapped_ref_);
}

void
Lorica::ReferenceMapper_i::remove_server(CORBA::Object_ptr mapped)
	throw (CORBA::SystemException)
{
	if (CORBA::is_nil(mapped))
		return;

	Lorica::ReferenceMapValue_var rmv = this->registry_->remove_reference (mapped, true);
	if ((rmv.get() != 0) && (rmv->ior_table_name_ != "")) {
		if (Lorica_debug_level > 5) {
			ACE_DEBUG((LM_DEBUG,
				   ACE_TEXT("(%P|%t) %N:%l - unbinding %s from IOR Table\n"),
				   rmv->ior_table_name_.c_str()));
		}

		this->ior_table_->unbind(rmv->ior_table_name_.c_str());
	} else {
		if (Lorica_debug_level > 5)
			ACE_DEBUG((LM_DEBUG,
				   ACE_TEXT("(%P|%t) %N:%l - remove reference returned rmv = 0x%x\n"),
				   rmv.get()));
	}

	// rmv releases reference to value which cleans itself up.
}

void
Lorica::ReferenceMapper_i::remove_client(CORBA::Object_ptr mapped)
	throw (CORBA::SystemException)
{
	if (CORBA::is_nil(mapped))
		return;

	Lorica::ReferenceMapValue_var rmv = this->registry_->remove_reference(mapped, false);
}

CORBA::Boolean
Lorica::ReferenceMapper_i::supports_security(void)
	throw (CORBA::SystemException)
{
	return this->has_security_;
}
<|endoftext|>