<commit_before>/*
 *  nextpnr -- Next Generation Place and Route
 *
 *  Copyright (C) 2019  David Shah <dave@ds0.me>
 *
 *  Permission to use, copy, modify, and/or distribute this software for any
 *  purpose with or without fee is hereby granted, provided that the above
 *  copyright notice and this permission notice appear in all copies.
 *
 *  THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
 *  WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
 *  MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
 *  ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
 *  WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
 *  ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
 *  OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 *
 */

#include "nextpnr.h"
#include "util.h"

NEXTPNR_NAMESPACE_BEGIN

namespace SDF {

struct MinMaxTyp
{
    double min, typ, max;
};

struct RiseFallDelay
{
    MinMaxTyp rise, fall;
};

struct PortAndEdge
{
    std::string port;
    ClockEdge edge;
};

struct IOPath
{
    std::string from, to;
    RiseFallDelay delay;
};

struct TimingCheck
{
    enum CheckType
    {
        SETUPHOLD,
        PERIOD,
        WIDTH
    } type;
    PortAndEdge from, to;
    RiseFallDelay delay;
};

struct Cell
{
    std::string celltype, instance;
    std::vector<IOPath> iopaths;
    std::vector<TimingCheck> checks;
};

struct CellPort
{
    std::string cell, port;
};

struct Interconnect
{
    CellPort from, to;
    RiseFallDelay delay;
};

struct SDFWriter
{
    std::vector<Cell> cells;
    std::vector<Interconnect> conn;
    std::string sdfversion, design, vendor, program;

    std::string format_name(const std::string &name)
    {
        std::string fmt = "\"";
        for (char c : name) {
            if (c == '\\' || c == '\"')
                fmt += "\"";
            fmt += c;
        }
        fmt += "\"";
        return fmt;
    }

    std::string escape_name(const std::string &name)
    {
        std::string esc;
        for (char c : name) {
            if (c == '$' || c == '\\' || c == '[' || c == ']' || c == ':')
                esc += '\\';
            esc += c;
        }
        return esc;
    }

    std::string timing_check_name(TimingCheck::CheckType type)
    {
        switch (type) {
        case TimingCheck::SETUPHOLD:
            return "SETUPHOLD";
        case TimingCheck::PERIOD:
            return "PERIOD";
        case TimingCheck::WIDTH:
            return "WIDTH";
        default:
            NPNR_ASSERT_FALSE("unknown timing check type");
        }
    }

    void write_delay(std::ostream &out, const RiseFallDelay &delay)
    {
        write_delay(out, delay.rise);
        out << " ";
        write_delay(out, delay.fall);
    }

    void write_delay(std::ostream &out, const MinMaxTyp &delay)
    {
        out << "(" << delay.min << ":" << delay.typ << ":" << delay.max << ")";
    }

    void write_port(std::ostream &out, const CellPort &port) { out << escape_name(port.cell + "/" + port.port); }

    void write_portedge(std::ostream &out, const PortAndEdge &pe)
    {
        out << "(" << (pe.edge == RISING_EDGE ? "posedge" : "negedge") << " " << escape_name(pe.port) << ")";
    }

    void write(std::ostream &out)
    {
        out << "(DELAYFILE" << std::endl;
        // Headers and  metadata
        out << "  (SDFVERSION " << format_name(sdfversion) << ")" << std::endl;
        out << "  (DESIGN " << format_name(design) << ")" << std::endl;
        out << "  (VENDOR " << format_name(vendor) << ")" << std::endl;
        out << "  (PROGRAM " << format_name(program) << ")" << std::endl;
        out << "  (DIVIDER /)" << std::endl;
        out << "  (TIMESCALE 1ps)" << std::endl;
        // Write interconnect delays, with the main design begin a "cell"
        out << "  (CELL" << std::endl;
        out << "    (CELLTYPE " << format_name(design) << ")" << std::endl;
        out << "    (INSTANCE )" << std::endl;
        out << "    (DELAY" << std::endl;
        out << "      (ABSOLUTE" << std::endl;
        for (auto &ic : conn) {
            out << "        (INTERCONNECT ";
            write_port(out, ic.from);
            out << " ";
            write_port(out, ic.to);
            out << " ";
            write_delay(out, ic.delay);
            out << ")" << std::endl;
        }
        out << "      )" << std::endl;
        out << "    )" << std::endl;
        out << "  )" << std::endl;
        // Write cells
        for (auto &cell : cells) {
            out << "  (CELL" << std::endl;
            out << "    (CELLTYPE " << format_name(cell.celltype) << ")" << std::endl;
            out << "    (INSTANCE " << escape_name(cell.instance) << ")" << std::endl;
            // IOPATHs (combinational delay and clock-to-q)
            if (!cell.iopaths.empty()) {
                out << "    (DELAY" << std::endl;
                out << "      (ABSOLUTE" << std::endl;
                for (auto &path : cell.iopaths) {
                    out << "        (IOPATH " << escape_name(path.from) << " " << escape_name(path.to) << " ";
                    write_delay(out, path.delay);
                    out << ")" << std::endl;
                }
                out << "      )" << std::endl;
                out << "    )" << std::endl;
            }
            // Timing Checks (setup/hold, period, width)
            if (!cell.checks.empty()) {
                out << "    (TIMINGCHECK" << std::endl;
                for (auto &check : cell.checks) {
                    out << "      (" << timing_check_name(check.type) << " ";
                    write_portedge(out, check.from);
                    out << " ";
                    if (check.type == TimingCheck::SETUPHOLD) {
                        write_portedge(out, check.to);
                        out << " ";
                    }
                    if (check.type == TimingCheck::SETUPHOLD)
                        write_delay(out, check.delay);
                    else
                        write_delay(out, check.delay.rise);
                    out << ")" << std::endl;
                }
                out << "    )" << std::endl;
            }
            out << "    )" << std::endl;
        }
        out << ")" << std::endl;
    }
};

} // namespace SDF

void Context::writeSDF(std::ostream &out) const
{
    using namespace SDF;
    SDFWriter wr;
    wr.design = str_or_default(attrs, id("module"), "top");
    wr.sdfversion = "3.0";
    wr.vendor = "nextpnr";
    wr.program = "nextpnr";

    const double delay_scale = 1000;
    // Convert from DelayInfo to SDF-friendly RiseFallDelay
    auto convert_delay = [&](const DelayInfo &dly) {
        RiseFallDelay rf;
        rf.rise.min = getDelayNS(dly.minRaiseDelay()) * delay_scale;
        rf.rise.typ = getDelayNS((dly.minRaiseDelay() + dly.maxRaiseDelay()) / 2) * delay_scale; // fixme: typ delays?
        rf.rise.max = getDelayNS(dly.maxRaiseDelay()) * delay_scale;
        rf.fall.min = getDelayNS(dly.minFallDelay()) * delay_scale;
        rf.fall.typ = getDelayNS((dly.minFallDelay() + dly.maxFallDelay()) / 2) * delay_scale; // fixme: typ delays?
        rf.fall.max = getDelayNS(dly.maxFallDelay()) * delay_scale;
        return rf;
    };

    auto convert_setuphold = [&](const DelayInfo &setup, const DelayInfo &hold) {
        RiseFallDelay rf;
        rf.rise.min = getDelayNS(setup.minDelay()) * delay_scale;
        rf.rise.typ = getDelayNS((setup.minDelay() + setup.maxDelay()) / 2) * delay_scale; // fixme: typ delays?
        rf.rise.max = getDelayNS(setup.maxDelay()) * delay_scale;
        rf.fall.min = getDelayNS(hold.minDelay()) * delay_scale;
        rf.fall.typ = getDelayNS((hold.minDelay() + hold.maxDelay()) / 2) * delay_scale; // fixme: typ delays?
        rf.fall.max = getDelayNS(hold.maxDelay()) * delay_scale;
        return rf;
    };

    for (auto cell : sorted(cells)) {
        Cell sc;
        const CellInfo *ci = cell.second;
        sc.instance = ci->name.str(this);
        sc.celltype = ci->type.str(this);
        for (auto port : ci->ports) {
            int clockCount = 0;
            TimingPortClass cls = getPortTimingClass(ci, port.first, clockCount);
            if (cls == TMG_IGNORE)
                continue;
            if (port.second.type != PORT_IN) {
                // Add combinational paths to this output (or inout)
                for (auto other : ci->ports) {
                    if (other.second.type == PORT_OUT)
                        continue;
                    DelayInfo dly;
                    if (!getCellDelay(ci, other.first, port.first, dly))
                        continue;
                    IOPath iop;
                    iop.from = other.first.str(this);
                    iop.to = port.first.str(this);
                    iop.delay = convert_delay(dly);
                    sc.iopaths.push_back(iop);
                }
                // Add clock-to-output delays, also as IOPaths
                if (cls == TMG_REGISTER_OUTPUT)
                    for (int i = 0; i < clockCount; i++) {
                        auto clkInfo = getPortClockingInfo(ci, port.first, i);
                        IOPath cqp;
                        cqp.from = clkInfo.clock_port.str(this);
                        cqp.to = port.first.str(this);
                        cqp.delay = convert_delay(clkInfo.clockToQ);
                        sc.iopaths.push_back(cqp);
                    }
            }
            if (port.second.type != PORT_OUT && cls == TMG_REGISTER_INPUT) {
                // Add setup/hold checks
                for (int i = 0; i < clockCount; i++) {
                    auto clkInfo = getPortClockingInfo(ci, port.first, i);
                    TimingCheck chk;
                    chk.from.edge = RISING_EDGE; // Add setup/hold checks equally for rising and falling edges
                    chk.from.port = port.first.str(this);
                    chk.to.edge = clkInfo.edge;
                    chk.to.port = clkInfo.clock_port.str(this);
                    chk.type = TimingCheck::SETUPHOLD;
                    chk.delay = convert_setuphold(clkInfo.setup, clkInfo.hold);
                    sc.checks.push_back(chk);
                    chk.from.edge = FALLING_EDGE;
                    sc.checks.push_back(chk);
                }
            }
        }
        wr.cells.push_back(sc);
    }

    for (auto net : sorted(nets)) {
        NetInfo *ni = net.second;
        if (ni->driver.cell == nullptr)
            continue;
        for (auto &usr : ni->users) {
            Interconnect ic;
            ic.from.cell = ni->driver.cell->name.str(this);
            ic.from.port = ni->driver.port.str(this);
            ic.to.cell = usr.cell->name.str(this);
            ic.to.port = usr.port.str(this);
            // FIXME: min/max routing delay - or at least constructing DelayInfo here
            ic.delay = convert_delay(getDelayFromNS(getDelayNS(getNetinfoRouteDelay(ni, usr))));
            wr.conn.push_back(ic);
        }
    }
    wr.write(out);
}

NEXTPNR_NAMESPACE_END<commit_msg>sdf: Improve SDF output<commit_after>/*
 *  nextpnr -- Next Generation Place and Route
 *
 *  Copyright (C) 2019  David Shah <dave@ds0.me>
 *
 *  Permission to use, copy, modify, and/or distribute this software for any
 *  purpose with or without fee is hereby granted, provided that the above
 *  copyright notice and this permission notice appear in all copies.
 *
 *  THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
 *  WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
 *  MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
 *  ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
 *  WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
 *  ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
 *  OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 *
 */

#include "nextpnr.h"
#include "util.h"

NEXTPNR_NAMESPACE_BEGIN

namespace SDF {

struct MinMaxTyp
{
    double min, typ, max;
};

struct RiseFallDelay
{
    MinMaxTyp rise, fall;
};

struct PortAndEdge
{
    std::string port;
    ClockEdge edge;
};

struct IOPath
{
    std::string from, to;
    RiseFallDelay delay;
};

struct TimingCheck
{
    enum CheckType
    {
        SETUPHOLD,
        PERIOD,
        WIDTH
    } type;
    PortAndEdge from, to;
    RiseFallDelay delay;
};

struct Cell
{
    std::string celltype, instance;
    std::vector<IOPath> iopaths;
    std::vector<TimingCheck> checks;
};

struct CellPort
{
    std::string cell, port;
};

struct Interconnect
{
    CellPort from, to;
    RiseFallDelay delay;
};

struct SDFWriter
{
    std::vector<Cell> cells;
    std::vector<Interconnect> conn;
    std::string sdfversion, design, vendor, program;

    std::string format_name(const std::string &name)
    {
        std::string fmt = "\"";
        for (char c : name) {
            if (c == '\\' || c == '\"')
                fmt += "\"";
            fmt += c;
        }
        fmt += "\"";
        return fmt;
    }

    std::string escape_name(const std::string &name)
    {
        std::string esc;
        for (char c : name) {
            if (c == '$' || c == '\\' || c == '[' || c == ']' || c == ':')
                esc += '\\';
            esc += c;
        }
        return esc;
    }

    std::string timing_check_name(TimingCheck::CheckType type)
    {
        switch (type) {
        case TimingCheck::SETUPHOLD:
            return "SETUPHOLD";
        case TimingCheck::PERIOD:
            return "PERIOD";
        case TimingCheck::WIDTH:
            return "WIDTH";
        default:
            NPNR_ASSERT_FALSE("unknown timing check type");
        }
    }

    void write_delay(std::ostream &out, const RiseFallDelay &delay)
    {
        write_delay(out, delay.rise);
        out << " ";
        write_delay(out, delay.fall);
    }

    void write_delay(std::ostream &out, const MinMaxTyp &delay)
    {
        out << "(" << delay.min << ":" << delay.typ << ":" << delay.max << ")";
    }

    void write_port(std::ostream &out, const CellPort &port) { out << escape_name(port.cell + "/" + port.port); }

    void write_portedge(std::ostream &out, const PortAndEdge &pe)
    {
        out << "(" << (pe.edge == RISING_EDGE ? "posedge" : "negedge") << " " << escape_name(pe.port) << ")";
    }

    void write(std::ostream &out)
    {
        out << "(DELAYFILE" << std::endl;
        // Headers and  metadata
        out << "  (SDFVERSION " << format_name(sdfversion) << ")" << std::endl;
        out << "  (DESIGN " << format_name(design) << ")" << std::endl;
        out << "  (VENDOR " << format_name(vendor) << ")" << std::endl;
        out << "  (PROGRAM " << format_name(program) << ")" << std::endl;
        out << "  (DIVIDER /)" << std::endl;
        out << "  (TIMESCALE 1ps)" << std::endl;
        // Write interconnect delays, with the main design begin a "cell"
        out << "  (CELL" << std::endl;
        out << "    (CELLTYPE " << format_name(design) << ")" << std::endl;
        out << "    (INSTANCE )" << std::endl;
        out << "    (DELAY" << std::endl;
        out << "      (ABSOLUTE" << std::endl;
        for (auto &ic : conn) {
            out << "        (INTERCONNECT ";
            write_port(out, ic.from);
            out << " ";
            write_port(out, ic.to);
            out << " ";
            write_delay(out, ic.delay);
            out << ")" << std::endl;
        }
        out << "      )" << std::endl;
        out << "    )" << std::endl;
        out << "  )" << std::endl;
        // Write cells
        for (auto &cell : cells) {
            out << "  (CELL" << std::endl;
            out << "    (CELLTYPE " << format_name(cell.celltype) << ")" << std::endl;
            out << "    (INSTANCE " << escape_name(cell.instance) << ")" << std::endl;
            // IOPATHs (combinational delay and clock-to-q)
            if (!cell.iopaths.empty()) {
                out << "    (DELAY" << std::endl;
                out << "      (ABSOLUTE" << std::endl;
                for (auto &path : cell.iopaths) {
                    out << "        (IOPATH " << escape_name(path.from) << " " << escape_name(path.to) << " ";
                    write_delay(out, path.delay);
                    out << ")" << std::endl;
                }
                out << "      )" << std::endl;
                out << "    )" << std::endl;
            }
            // Timing Checks (setup/hold, period, width)
            if (!cell.checks.empty()) {
                out << "    (TIMINGCHECK" << std::endl;
                for (auto &check : cell.checks) {
                    out << "      (" << timing_check_name(check.type) << " ";
                    write_portedge(out, check.from);
                    out << " ";
                    if (check.type == TimingCheck::SETUPHOLD) {
                        write_portedge(out, check.to);
                        out << " ";
                    }
                    if (check.type == TimingCheck::SETUPHOLD)
                        write_delay(out, check.delay);
                    else
                        write_delay(out, check.delay.rise);
                    out << ")" << std::endl;
                }
                out << "    )" << std::endl;
            }
            out << "    )" << std::endl;
        }
        out << ")" << std::endl;
    }
};

} // namespace SDF

void Context::writeSDF(std::ostream &out) const
{
    using namespace SDF;
    SDFWriter wr;
    wr.design = str_or_default(attrs, id("module"), "top");
    wr.sdfversion = "3.0";
    wr.vendor = "nextpnr";
    wr.program = "nextpnr";

    const double delay_scale = 1000;
    // Convert from DelayInfo to SDF-friendly RiseFallDelay
    auto convert_delay = [&](const DelayInfo &dly) {
        RiseFallDelay rf;
        rf.rise.min = getDelayNS(dly.minRaiseDelay()) * delay_scale;
        rf.rise.typ = getDelayNS((dly.minRaiseDelay() + dly.maxRaiseDelay()) / 2) * delay_scale; // fixme: typ delays?
        rf.rise.max = getDelayNS(dly.maxRaiseDelay()) * delay_scale;
        rf.fall.min = getDelayNS(dly.minFallDelay()) * delay_scale;
        rf.fall.typ = getDelayNS((dly.minFallDelay() + dly.maxFallDelay()) / 2) * delay_scale; // fixme: typ delays?
        rf.fall.max = getDelayNS(dly.maxFallDelay()) * delay_scale;
        return rf;
    };

    auto convert_setuphold = [&](const DelayInfo &setup, const DelayInfo &hold) {
        RiseFallDelay rf;
        rf.rise.min = getDelayNS(setup.minDelay()) * delay_scale;
        rf.rise.typ = getDelayNS((setup.minDelay() + setup.maxDelay()) / 2) * delay_scale; // fixme: typ delays?
        rf.rise.max = getDelayNS(setup.maxDelay()) * delay_scale;
        rf.fall.min = getDelayNS(hold.minDelay()) * delay_scale;
        rf.fall.typ = getDelayNS((hold.minDelay() + hold.maxDelay()) / 2) * delay_scale; // fixme: typ delays?
        rf.fall.max = getDelayNS(hold.maxDelay()) * delay_scale;
        return rf;
    };

    for (auto cell : sorted(cells)) {
        Cell sc;
        const CellInfo *ci = cell.second;
        sc.instance = ci->name.str(this);
        sc.celltype = ci->type.str(this);
        for (auto port : ci->ports) {
            int clockCount = 0;
            TimingPortClass cls = getPortTimingClass(ci, port.first, clockCount);
            if (cls == TMG_IGNORE)
                continue;
            if (port.second.net == nullptr)
                continue; // Ignore disconnected ports
            if (port.second.type != PORT_IN) {
                // Add combinational paths to this output (or inout)
                for (auto other : ci->ports) {
                    if (other.second.net == nullptr)
                        continue;
                    if (other.second.type == PORT_OUT)
                        continue;
                    DelayInfo dly;
                    if (!getCellDelay(ci, other.first, port.first, dly))
                        continue;
                    IOPath iop;
                    iop.from = other.first.str(this);
                    iop.to = port.first.str(this);
                    iop.delay = convert_delay(dly);
                    sc.iopaths.push_back(iop);
                }
                // Add clock-to-output delays, also as IOPaths
                if (cls == TMG_REGISTER_OUTPUT)
                    for (int i = 0; i < clockCount; i++) {
                        auto clkInfo = getPortClockingInfo(ci, port.first, i);
                        IOPath cqp;
                        cqp.from = clkInfo.clock_port.str(this);
                        cqp.to = port.first.str(this);
                        cqp.delay = convert_delay(clkInfo.clockToQ);
                        sc.iopaths.push_back(cqp);
                    }
            }
            if (port.second.type != PORT_OUT && cls == TMG_REGISTER_INPUT) {
                // Add setup/hold checks
                for (int i = 0; i < clockCount; i++) {
                    auto clkInfo = getPortClockingInfo(ci, port.first, i);
                    TimingCheck chk;
                    chk.from.edge = RISING_EDGE; // Add setup/hold checks equally for rising and falling edges
                    chk.from.port = port.first.str(this);
                    chk.to.edge = clkInfo.edge;
                    chk.to.port = clkInfo.clock_port.str(this);
                    chk.type = TimingCheck::SETUPHOLD;
                    chk.delay = convert_setuphold(clkInfo.setup, clkInfo.hold);
                    sc.checks.push_back(chk);
                    chk.from.edge = FALLING_EDGE;
                    sc.checks.push_back(chk);
                }
            }
        }
        wr.cells.push_back(sc);
    }

    for (auto net : sorted(nets)) {
        NetInfo *ni = net.second;
        if (ni->driver.cell == nullptr)
            continue;
        for (auto &usr : ni->users) {
            Interconnect ic;
            ic.from.cell = ni->driver.cell->name.str(this);
            ic.from.port = ni->driver.port.str(this);
            ic.to.cell = usr.cell->name.str(this);
            ic.to.port = usr.port.str(this);
            // FIXME: min/max routing delay - or at least constructing DelayInfo here
            ic.delay = convert_delay(getDelayFromNS(getDelayNS(getNetinfoRouteDelay(ni, usr))));
            wr.conn.push_back(ic);
        }
    }
    wr.write(out);
}

NEXTPNR_NAMESPACE_END<|endoftext|>


<commit_before>#include "AMDGraph.h"

namespace pf
{
namespace graph
{

template<typename K, typename V>
AMDGraph<K, V>::AMDGraph()
{

}

template<typename K, typename V>
AMDGraph<K, V>& AMDGraph<K, V>::addVertex(const K& v)
{
	vl.push_back(v);

	am.reserve(vl.size());
	std::for_each(am.begin(), am.end(), [&](std::vector<V>& v){v.reserve(vl.size());});

	for(int i = 0; i < am.size() - 1; ++i)
	{
		am[vl.size() - 1][i] = nullptr;
		am[i][vl.size() - 1] = nullptr;
	}
	am[vl.size() - 1][vl.size() - 1] = nullptr;

	return this;
}

template<typename K, typename V>
bool AMDGraph<K, V>::addEdge(const K& v1, const K& v2, const V* weight)
{
	std::vector<K>::iterator v1It = std::find(vl.begin(), vl.end(), v1);
	std::vector<K>::iterator v2It = std::find(vl.begin(), vl.end(), v2);

	if(v1It != vl.end() && v2It != vl.end())
	{
		*((am.begin() + v1It)->begin() + v2It) = weight;
		return true;
	}
	return false;
}

template<typename K, typename V>
bool AMDGraph<K, V>::removeVertex(const K& v)
{
	std::vector<K>::iterator it = std::find(vl.begin(), vl.end(), v);

	if(it != vl.end())
	{
		vl.erase(it);

		am.erase(am.begin() + it);

		for(int i = 0; i < am.size(); ++i)
		{
			am[i].erase(am[i].begin() + it);
		}

		return true;
	}
	return false;
}

template<typename K, typename V>
bool AMDGraph<K, V>::removeEdge(const K& v1, const K& v2)
{
	std::vector<K>::iterator v1It = std::find(vl.begin(), vl.end(), v1);
	std::vector<K>::iterator v2It = std::find(vl.begin(), vl.end(), v2);

	if(v1It != vl.end() && v2It != vl.end())
	{
		*((am.begin() + v1It)->begin() + v2It) = nullptr;
		return true;
	}
	return false;
}

template<typename K, typename V>
AMDGraph<K, V>::~AMDGraph()
{
	~vl();
	for_each(am.begin(), am.end(), [](std::vector<V>& v){~v();});
	~am();
}

template<typename K, typename V>
std::ostream& operator<<(std::ostream& out, AMDGraph<K, V> graph)
{
	//todo do the needfull
	//todo make friend of AMDGraph?
}

}
}
<commit_msg>AMDGraph needs template specializations for bool (unweighed graphs)<commit_after>#include "AMDGraph.h"

namespace pf
{
namespace graph
{

template<typename K, typename V>
AMDGraph<K, V>::AMDGraph()
{

}

template<typename K, typename V>
AMDGraph<K, V>& AMDGraph<K, V>::addVertex(const K& v)
{
	vl.push_back(v);

	am.reserve(vl.size());
	std::for_each(am.begin(), am.end(), [&](std::vector<V>& v){v.reserve(vl.size());});

	for(int i = 0; i < am.size() - 1; ++i)
	{
		am[vl.size() - 1][i] = nullptr;
		am[i][vl.size() - 1] = nullptr;
	}
	am[vl.size() - 1][vl.size() - 1] = nullptr;

	return this;
}

template<typename K, typename V>
bool AMDGraph<K, V>::addEdge(const K& v1, const K& v2, const V* weight)
{
	std::vector<K>::iterator v1It = std::find(vl.begin(), vl.end(), v1);
	std::vector<K>::iterator v2It = std::find(vl.begin(), vl.end(), v2);

	if(v1It != vl.end() && v2It != vl.end())
	{
		*((am.begin() + v1It)->begin() + v2It) = weight;
		return true;
	}
	return false;
}

/*
template<typename K>
bool AMDGraph<K, bool>::addEdge(const K& v1, const K& v2, const bool* weight)
{
	std::vector<K>::iterator v1It = std::find(vl.begin(), vl.end(), v1);
	std::vector<K>::iterator v2It = std::find(vl.begin(), vl.end(), v2);

	if(v1It != vl.end() && v2It != vl.end())
	{
		*((am.begin() + v1It)->begin() + v2It) = weight;
		return true;
	}
	return false;
}
*/

template<typename K, typename V>
bool AMDGraph<K, V>::removeVertex(const K& v)
{
	std::vector<K>::iterator it = std::find(vl.begin(), vl.end(), v);

	if(it != vl.end())
	{
		vl.erase(it);

		am.erase(am.begin() + it);

		for(int i = 0; i < am.size(); ++i)
		{
			am[i].erase(am[i].begin() + it);
		}

		return true;
	}
	return false;
}

template<typename K, typename V>
bool AMDGraph<K, V>::removeEdge(const K& v1, const K& v2)
{
	std::vector<K>::iterator v1It = std::find(vl.begin(), vl.end(), v1);
	std::vector<K>::iterator v2It = std::find(vl.begin(), vl.end(), v2);

	if(v1It != vl.end() && v2It != vl.end())
	{
		*((am.begin() + v1It)->begin() + v2It) = nullptr;
		return true;
	}
	return false;
}

template<typename K, typename V>
AMDGraph<K, V>::~AMDGraph()
{
	~vl();
	for_each(am.begin(), am.end(), [](std::vector<V>& v){~v();});
	~am();
}

template<typename K, typename V>
std::ostream& operator<<(std::ostream& out, AMDGraph<K, V> graph)
{
	//todo do the needfull
	//todo make friend of AMDGraph?
}

}
}
<|endoftext|>


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

#include "riegeli/bytes/string_writer.h"

#include <stddef.h>

#include <string>
#include <utility>

#include "absl/base/optimization.h"
#include "absl/strings/cord.h"
#include "absl/strings/string_view.h"
#include "absl/types/optional.h"
#include "absl/types/span.h"
#include "riegeli/base/base.h"
#include "riegeli/base/chain.h"
#include "riegeli/bytes/reader.h"
#include "riegeli/bytes/string_reader.h"
#include "riegeli/bytes/writer.h"

namespace riegeli {

void StringWriterBase::Done() {
  StringWriterBase::FlushImpl(FlushType::kFromObject);
  Writer::Done();
  secondary_buffer_ = Chain();
  associated_reader_.Reset();
}

bool StringWriterBase::PushSlow(size_t min_length, size_t recommended_length) {
  RIEGELI_ASSERT_LT(available(), min_length)
      << "Failed precondition of Writer::PushSlow(): "
         "enough space available, use Push() instead";
  if (ABSL_PREDICT_FALSE(!ok())) return false;
  std::string& dest = *dest_string();
  RIEGELI_ASSERT_EQ(limit_pos(), dest.size() + secondary_buffer_.size())
      << "StringWriter destination changed unexpectedly";
  if (ABSL_PREDICT_FALSE(min_length >
                         dest.max_size() - IntCast<size_t>(pos()))) {
    return FailOverflow();
  }
  if (secondary_buffer_.empty()) {
    SyncDestBuffer(dest);
    if (dest.size() <= std::string().capacity()) {
      // Allocate the first block directly in `dest`. It is possible that it
      // will not need to be copied if it turns out to be the only block,
      // although this decision might cause it to remain wasteful if less data
      // are written than space requested.
      dest.reserve(UnsignedMax(
          dest.size() + UnsignedMax(min_length, recommended_length),
          // Subtract 1 for NUL terminator because `min_block_size()` might be
          // a good fit for an allocation size.
          SaturatingSub(options_.min_block_size(), size_t{1})));
    }
    if (min_length <= dest.capacity() - dest.size()) {
      MakeDestBuffer(dest);
      return true;
    }
    set_start_pos(dest.size());
  } else {
    SyncSecondaryBuffer();
  }
  MakeSecondaryBuffer(min_length, recommended_length);
  return true;
}

bool StringWriterBase::WriteSlow(const Chain& src) {
  RIEGELI_ASSERT_LT(UnsignedMin(available(), kMaxBytesToCopy), src.size())
      << "Failed precondition of Writer::WriteSlow(Chain): "
         "enough space available, use Write(Chain) instead";
  if (ABSL_PREDICT_FALSE(!ok())) return false;
  std::string& dest = *dest_string();
  RIEGELI_ASSERT_EQ(limit_pos(), dest.size() + secondary_buffer_.size())
      << "StringWriter destination changed unexpectedly";
  if (ABSL_PREDICT_FALSE(src.size() >
                         dest.max_size() - IntCast<size_t>(pos()))) {
    return FailOverflow();
  }
  if (secondary_buffer_.empty()) {
    SyncDestBuffer(dest);
    if (src.size() <= dest.capacity() - dest.size()) {
      src.AppendTo(dest);
      MakeDestBuffer(dest);
      return true;
    }
    set_start_pos(dest.size());
  } else {
    SyncSecondaryBuffer();
  }
  move_start_pos(src.size());
  secondary_buffer_.Append(src, options_);
  MakeSecondaryBuffer();
  return true;
}

bool StringWriterBase::WriteSlow(Chain&& src) {
  RIEGELI_ASSERT_LT(UnsignedMin(available(), kMaxBytesToCopy), src.size())
      << "Failed precondition of Writer::WriteSlow(Chain&&): "
         "enough space available, use Write(Chain) instead";
  if (ABSL_PREDICT_FALSE(!ok())) return false;
  std::string& dest = *dest_string();
  RIEGELI_ASSERT_EQ(limit_pos(), dest.size() + secondary_buffer_.size())
      << "StringWriter destination changed unexpectedly";
  if (ABSL_PREDICT_FALSE(src.size() >
                         dest.max_size() - IntCast<size_t>(pos()))) {
    return FailOverflow();
  }
  if (secondary_buffer_.empty()) {
    SyncDestBuffer(dest);
    if (src.size() <= dest.capacity() - dest.size()) {
      std::move(src).AppendTo(dest);
      MakeDestBuffer(dest);
      return true;
    }
    set_start_pos(dest.size());
  } else {
    SyncSecondaryBuffer();
  }
  move_start_pos(src.size());
  secondary_buffer_.Append(std::move(src), options_);
  MakeSecondaryBuffer();
  return true;
}

bool StringWriterBase::WriteSlow(const absl::Cord& src) {
  RIEGELI_ASSERT_LT(UnsignedMin(available(), kMaxBytesToCopy), src.size())
      << "Failed precondition of Writer::WriteSlow(Cord): "
         "enough space available, use Write(Cord) instead";
  if (ABSL_PREDICT_FALSE(!ok())) return false;
  std::string& dest = *dest_string();
  RIEGELI_ASSERT_EQ(limit_pos(), dest.size() + secondary_buffer_.size())
      << "StringWriter destination changed unexpectedly";
  if (ABSL_PREDICT_FALSE(src.size() >
                         dest.max_size() - IntCast<size_t>(pos()))) {
    return FailOverflow();
  }
  if (secondary_buffer_.empty()) {
    SyncDestBuffer(dest);
    if (src.size() <= dest.capacity() - dest.size()) {
      for (const absl::string_view fragment : src.Chunks()) {
        // TODO: When `absl::string_view` becomes C++17
        // `std::string_view`: `dest.append(fragment)`
        dest.append(fragment.data(), fragment.size());
      }
      MakeDestBuffer(dest);
      return true;
    }
    set_start_pos(dest.size());
  } else {
    SyncSecondaryBuffer();
  }
  move_start_pos(src.size());
  secondary_buffer_.Append(src, options_);
  MakeSecondaryBuffer();
  return true;
}

bool StringWriterBase::WriteSlow(absl::Cord&& src) {
  RIEGELI_ASSERT_LT(UnsignedMin(available(), kMaxBytesToCopy), src.size())
      << "Failed precondition of Writer::WriteSlow(Cord&&): "
         "enough space available, use Write(Cord&&) instead";
  if (ABSL_PREDICT_FALSE(!ok())) return false;
  std::string& dest = *dest_string();
  RIEGELI_ASSERT_EQ(limit_pos(), dest.size() + secondary_buffer_.size())
      << "StringWriter destination changed unexpectedly";
  if (ABSL_PREDICT_FALSE(src.size() >
                         dest.max_size() - IntCast<size_t>(pos()))) {
    return FailOverflow();
  }
  if (secondary_buffer_.empty()) {
    SyncDestBuffer(dest);
    if (src.size() <= dest.capacity() - dest.size()) {
      for (const absl::string_view fragment : src.Chunks()) {
        // TODO: When `absl::string_view` becomes C++17
        // `std::string_view`: `dest.append(fragment)`
        dest.append(fragment.data(), fragment.size());
      }
      MakeDestBuffer(dest);
      return true;
    }
    set_start_pos(dest.size());
  } else {
    SyncSecondaryBuffer();
  }
  move_start_pos(src.size());
  secondary_buffer_.Append(std::move(src), options_);
  MakeSecondaryBuffer();
  return true;
}

bool StringWriterBase::WriteZerosSlow(Position length) {
  RIEGELI_ASSERT_LT(UnsignedMin(available(), kMaxBytesToCopy), length)
      << "Failed precondition of Writer::WriteZerosSlow(): "
         "enough space available, use WriteZeros() instead";
  if (ABSL_PREDICT_FALSE(!ok())) return false;
  std::string& dest = *dest_string();
  RIEGELI_ASSERT_EQ(limit_pos(), dest.size() + secondary_buffer_.size())
      << "StringWriter destination changed unexpectedly";
  if (ABSL_PREDICT_FALSE(length > dest.max_size() - IntCast<size_t>(pos()))) {
    return FailOverflow();
  }
  if (secondary_buffer_.empty()) {
    SyncDestBuffer(dest);
    if (length <= dest.capacity() - dest.size()) {
      dest.append(length, '\0');
      MakeDestBuffer(dest);
      return true;
    }
    set_start_pos(dest.size());
  } else {
    SyncSecondaryBuffer();
  }
  move_start_pos(length);
  secondary_buffer_.Append(ChainOfZeros(IntCast<size_t>(length)), options_);
  MakeSecondaryBuffer();
  return true;
}

bool StringWriterBase::FlushImpl(FlushType flush_type) {
  if (ABSL_PREDICT_FALSE(!ok())) return false;
  std::string& dest = *dest_string();
  RIEGELI_ASSERT_EQ(limit_pos(), dest.size() + secondary_buffer_.size())
      << "StringWriter destination changed unexpectedly";
  if (secondary_buffer_.empty()) {
    SyncDestBuffer(dest);
  } else {
    SyncSecondaryBuffer();
    std::move(secondary_buffer_).AppendTo(dest);
    secondary_buffer_.Clear();
    set_start_pos(0);
    set_buffer(&dest[0], dest.size(), dest.size());
  }
  return true;
}

absl::optional<Position> StringWriterBase::SizeImpl() {
  if (ABSL_PREDICT_FALSE(!ok())) return absl::nullopt;
  return pos();
}

bool StringWriterBase::TruncateImpl(Position new_size) {
  if (ABSL_PREDICT_FALSE(!ok())) return false;
  std::string& dest = *dest_string();
  RIEGELI_ASSERT_EQ(limit_pos(), dest.size() + secondary_buffer_.size())
      << "StringWriter destination changed unexpectedly";
  if (ABSL_PREDICT_FALSE(new_size > pos())) return false;
  if (new_size < dest.size()) {
    secondary_buffer_.Clear();
    set_start_pos(0);
    set_buffer(&dest[0], dest.size(), IntCast<size_t>(new_size));
  } else {
    secondary_buffer_.RemoveSuffix(
        dest.size() + secondary_buffer_.size() - IntCast<size_t>(new_size),
        options_);
    set_start_pos(new_size);
    set_buffer();
  }
  return true;
}

Reader* StringWriterBase::ReadModeImpl(Position initial_pos) {
  if (ABSL_PREDICT_FALSE(
          !StringWriterBase::FlushImpl(FlushType::kFromObject))) {
    return nullptr;
  }
  std::string& dest = *dest_string();
  StringReader<>* const reader = associated_reader_.ResetReader(dest);
  reader->Seek(initial_pos);
  return reader;
}

inline void StringWriterBase::SyncDestBuffer(std::string& dest) {
  dest.erase(dest.size() - available());
  set_buffer(&dest[0], dest.size(), dest.size());
}

inline void StringWriterBase::SyncSecondaryBuffer() {
  set_start_pos(pos());
  secondary_buffer_.RemoveSuffix(available(), options_);
  set_buffer();
}

inline void StringWriterBase::MakeSecondaryBuffer(size_t min_length,
                                                  size_t recommended_length) {
  const absl::Span<char> buffer = secondary_buffer_.AppendBuffer(
      min_length, recommended_length, Chain::kAnyLength, options_);
  set_buffer(buffer.data(), buffer.size());
}

}  // namespace riegeli
<commit_msg>In `StringWriter`, reduce the capacity of a smallish `std::string` (affecting sizes between 23 and 255 bytes on Linux with modern libc++ and with default `StringWriterBase::Options`).<commit_after>// Copyright 2017 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include "riegeli/bytes/string_writer.h"

#include <stddef.h>

#include <string>
#include <utility>

#include "absl/base/optimization.h"
#include "absl/strings/cord.h"
#include "absl/strings/string_view.h"
#include "absl/types/optional.h"
#include "absl/types/span.h"
#include "riegeli/base/base.h"
#include "riegeli/base/chain.h"
#include "riegeli/bytes/reader.h"
#include "riegeli/bytes/string_reader.h"
#include "riegeli/bytes/writer.h"

namespace riegeli {

void StringWriterBase::Done() {
  StringWriterBase::FlushImpl(FlushType::kFromObject);
  Writer::Done();
  secondary_buffer_ = Chain();
  associated_reader_.Reset();
}

bool StringWriterBase::PushSlow(size_t min_length, size_t recommended_length) {
  RIEGELI_ASSERT_LT(available(), min_length)
      << "Failed precondition of Writer::PushSlow(): "
         "enough space available, use Push() instead";
  if (ABSL_PREDICT_FALSE(!ok())) return false;
  std::string& dest = *dest_string();
  RIEGELI_ASSERT_EQ(limit_pos(), dest.size() + secondary_buffer_.size())
      << "StringWriter destination changed unexpectedly";
  if (ABSL_PREDICT_FALSE(min_length >
                         dest.max_size() - IntCast<size_t>(pos()))) {
    return FailOverflow();
  }
  if (secondary_buffer_.empty()) {
    SyncDestBuffer(dest);
    if (dest.size() <= std::string().capacity()) {
      // Allocate the first block directly in `dest`. It is possible that it
      // will not need to be copied if it turns out to be the only block,
      // although this decision might cause it to remain wasteful if less data
      // are written than space requested.
      dest.reserve(dest.size() + UnsignedMax(min_length, recommended_length,
                                             dest.size() / 2));
    }
    if (min_length <= dest.capacity() - dest.size()) {
      MakeDestBuffer(dest);
      return true;
    }
    set_start_pos(dest.size());
  } else {
    SyncSecondaryBuffer();
  }
  MakeSecondaryBuffer(min_length, recommended_length);
  return true;
}

bool StringWriterBase::WriteSlow(const Chain& src) {
  RIEGELI_ASSERT_LT(UnsignedMin(available(), kMaxBytesToCopy), src.size())
      << "Failed precondition of Writer::WriteSlow(Chain): "
         "enough space available, use Write(Chain) instead";
  if (ABSL_PREDICT_FALSE(!ok())) return false;
  std::string& dest = *dest_string();
  RIEGELI_ASSERT_EQ(limit_pos(), dest.size() + secondary_buffer_.size())
      << "StringWriter destination changed unexpectedly";
  if (ABSL_PREDICT_FALSE(src.size() >
                         dest.max_size() - IntCast<size_t>(pos()))) {
    return FailOverflow();
  }
  if (secondary_buffer_.empty()) {
    SyncDestBuffer(dest);
    if (src.size() <= dest.capacity() - dest.size()) {
      src.AppendTo(dest);
      MakeDestBuffer(dest);
      return true;
    }
    set_start_pos(dest.size());
  } else {
    SyncSecondaryBuffer();
  }
  move_start_pos(src.size());
  secondary_buffer_.Append(src, options_);
  MakeSecondaryBuffer();
  return true;
}

bool StringWriterBase::WriteSlow(Chain&& src) {
  RIEGELI_ASSERT_LT(UnsignedMin(available(), kMaxBytesToCopy), src.size())
      << "Failed precondition of Writer::WriteSlow(Chain&&): "
         "enough space available, use Write(Chain) instead";
  if (ABSL_PREDICT_FALSE(!ok())) return false;
  std::string& dest = *dest_string();
  RIEGELI_ASSERT_EQ(limit_pos(), dest.size() + secondary_buffer_.size())
      << "StringWriter destination changed unexpectedly";
  if (ABSL_PREDICT_FALSE(src.size() >
                         dest.max_size() - IntCast<size_t>(pos()))) {
    return FailOverflow();
  }
  if (secondary_buffer_.empty()) {
    SyncDestBuffer(dest);
    if (src.size() <= dest.capacity() - dest.size()) {
      std::move(src).AppendTo(dest);
      MakeDestBuffer(dest);
      return true;
    }
    set_start_pos(dest.size());
  } else {
    SyncSecondaryBuffer();
  }
  move_start_pos(src.size());
  secondary_buffer_.Append(std::move(src), options_);
  MakeSecondaryBuffer();
  return true;
}

bool StringWriterBase::WriteSlow(const absl::Cord& src) {
  RIEGELI_ASSERT_LT(UnsignedMin(available(), kMaxBytesToCopy), src.size())
      << "Failed precondition of Writer::WriteSlow(Cord): "
         "enough space available, use Write(Cord) instead";
  if (ABSL_PREDICT_FALSE(!ok())) return false;
  std::string& dest = *dest_string();
  RIEGELI_ASSERT_EQ(limit_pos(), dest.size() + secondary_buffer_.size())
      << "StringWriter destination changed unexpectedly";
  if (ABSL_PREDICT_FALSE(src.size() >
                         dest.max_size() - IntCast<size_t>(pos()))) {
    return FailOverflow();
  }
  if (secondary_buffer_.empty()) {
    SyncDestBuffer(dest);
    if (src.size() <= dest.capacity() - dest.size()) {
      for (const absl::string_view fragment : src.Chunks()) {
        // TODO: When `absl::string_view` becomes C++17
        // `std::string_view`: `dest.append(fragment)`
        dest.append(fragment.data(), fragment.size());
      }
      MakeDestBuffer(dest);
      return true;
    }
    set_start_pos(dest.size());
  } else {
    SyncSecondaryBuffer();
  }
  move_start_pos(src.size());
  secondary_buffer_.Append(src, options_);
  MakeSecondaryBuffer();
  return true;
}

bool StringWriterBase::WriteSlow(absl::Cord&& src) {
  RIEGELI_ASSERT_LT(UnsignedMin(available(), kMaxBytesToCopy), src.size())
      << "Failed precondition of Writer::WriteSlow(Cord&&): "
         "enough space available, use Write(Cord&&) instead";
  if (ABSL_PREDICT_FALSE(!ok())) return false;
  std::string& dest = *dest_string();
  RIEGELI_ASSERT_EQ(limit_pos(), dest.size() + secondary_buffer_.size())
      << "StringWriter destination changed unexpectedly";
  if (ABSL_PREDICT_FALSE(src.size() >
                         dest.max_size() - IntCast<size_t>(pos()))) {
    return FailOverflow();
  }
  if (secondary_buffer_.empty()) {
    SyncDestBuffer(dest);
    if (src.size() <= dest.capacity() - dest.size()) {
      for (const absl::string_view fragment : src.Chunks()) {
        // TODO: When `absl::string_view` becomes C++17
        // `std::string_view`: `dest.append(fragment)`
        dest.append(fragment.data(), fragment.size());
      }
      MakeDestBuffer(dest);
      return true;
    }
    set_start_pos(dest.size());
  } else {
    SyncSecondaryBuffer();
  }
  move_start_pos(src.size());
  secondary_buffer_.Append(std::move(src), options_);
  MakeSecondaryBuffer();
  return true;
}

bool StringWriterBase::WriteZerosSlow(Position length) {
  RIEGELI_ASSERT_LT(UnsignedMin(available(), kMaxBytesToCopy), length)
      << "Failed precondition of Writer::WriteZerosSlow(): "
         "enough space available, use WriteZeros() instead";
  if (ABSL_PREDICT_FALSE(!ok())) return false;
  std::string& dest = *dest_string();
  RIEGELI_ASSERT_EQ(limit_pos(), dest.size() + secondary_buffer_.size())
      << "StringWriter destination changed unexpectedly";
  if (ABSL_PREDICT_FALSE(length > dest.max_size() - IntCast<size_t>(pos()))) {
    return FailOverflow();
  }
  if (secondary_buffer_.empty()) {
    SyncDestBuffer(dest);
    if (length <= dest.capacity() - dest.size()) {
      dest.append(length, '\0');
      MakeDestBuffer(dest);
      return true;
    }
    set_start_pos(dest.size());
  } else {
    SyncSecondaryBuffer();
  }
  move_start_pos(length);
  secondary_buffer_.Append(ChainOfZeros(IntCast<size_t>(length)), options_);
  MakeSecondaryBuffer();
  return true;
}

bool StringWriterBase::FlushImpl(FlushType flush_type) {
  if (ABSL_PREDICT_FALSE(!ok())) return false;
  std::string& dest = *dest_string();
  RIEGELI_ASSERT_EQ(limit_pos(), dest.size() + secondary_buffer_.size())
      << "StringWriter destination changed unexpectedly";
  if (secondary_buffer_.empty()) {
    SyncDestBuffer(dest);
  } else {
    SyncSecondaryBuffer();
    std::move(secondary_buffer_).AppendTo(dest);
    secondary_buffer_.Clear();
    set_start_pos(0);
    set_buffer(&dest[0], dest.size(), dest.size());
  }
  return true;
}

absl::optional<Position> StringWriterBase::SizeImpl() {
  if (ABSL_PREDICT_FALSE(!ok())) return absl::nullopt;
  return pos();
}

bool StringWriterBase::TruncateImpl(Position new_size) {
  if (ABSL_PREDICT_FALSE(!ok())) return false;
  std::string& dest = *dest_string();
  RIEGELI_ASSERT_EQ(limit_pos(), dest.size() + secondary_buffer_.size())
      << "StringWriter destination changed unexpectedly";
  if (ABSL_PREDICT_FALSE(new_size > pos())) return false;
  if (new_size < dest.size()) {
    secondary_buffer_.Clear();
    set_start_pos(0);
    set_buffer(&dest[0], dest.size(), IntCast<size_t>(new_size));
  } else {
    secondary_buffer_.RemoveSuffix(
        dest.size() + secondary_buffer_.size() - IntCast<size_t>(new_size),
        options_);
    set_start_pos(new_size);
    set_buffer();
  }
  return true;
}

Reader* StringWriterBase::ReadModeImpl(Position initial_pos) {
  if (ABSL_PREDICT_FALSE(
          !StringWriterBase::FlushImpl(FlushType::kFromObject))) {
    return nullptr;
  }
  std::string& dest = *dest_string();
  StringReader<>* const reader = associated_reader_.ResetReader(dest);
  reader->Seek(initial_pos);
  return reader;
}

inline void StringWriterBase::SyncDestBuffer(std::string& dest) {
  dest.erase(dest.size() - available());
  set_buffer(&dest[0], dest.size(), dest.size());
}

inline void StringWriterBase::SyncSecondaryBuffer() {
  set_start_pos(pos());
  secondary_buffer_.RemoveSuffix(available(), options_);
  set_buffer();
}

inline void StringWriterBase::MakeSecondaryBuffer(size_t min_length,
                                                  size_t recommended_length) {
  const absl::Span<char> buffer = secondary_buffer_.AppendBuffer(
      min_length, recommended_length, Chain::kAnyLength, options_);
  set_buffer(buffer.data(), buffer.size());
}

}  // namespace riegeli
<|endoftext|>


<commit_before><commit_msg>p1018<commit_after><|endoftext|>


<commit_before><commit_msg>pdn: ensure instance grids only add their obs obstructions and not the aggregate of all the prior grids<commit_after><|endoftext|>


<commit_before>#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include <Eigen/Dense>
#include <unsupported/Eigen/AutoDiff>

#include <bundle/src/projection_errors.h>

class ReprojectionError2DFixture : public ::testing::Test {
 public:
  typedef Eigen::AutoDiffScalar<Eigen::VectorXd> AScalar;
  ReprojectionError2DFixture() { observed << 0.5, 0.5; }

  void SetupADiff(int size, const double* camera, AScalar* camera_adiff) {
    const int total_size = size_point + size_rt + size;
    for (int i = 0; i < size_point; ++i) {
      point_adiff[i].value() = point[i];
      point_adiff[i].derivatives() = VecXd::Unit(total_size, i);
    }
    for (int i = 0; i < size_rt; ++i) {
      rt_adiff[i].value() = rt[i];
      rt_adiff[i].derivatives() = VecXd::Unit(total_size, size_point + i);
    }
    for (int i = 0; i < size; ++i) {
      camera_adiff[i].value() = camera[i];
      camera_adiff[i].derivatives() = VecXd::Unit(total_size, size_point + size_rt + i);
    }
  }

  void CheckJacobians(int size, const double* jac_camera) {
    for (int i = 0; i < size_residual; ++i) {
      for (int j = 0; j < size; ++j) {
        ASSERT_NEAR(residual_adiff[i].derivatives()(size_point + size_rt + j),
                    jac_camera[i * size + j], 1e-15);
      }
    }
    for (int i = 0; i < size_residual; ++i) {
      for (int j = 0; j < size_point; ++j) {
        ASSERT_NEAR(residual_adiff[i].derivatives()(j),
                    jac_point[i * size_point + j], 1e-15);
      }
    }
    for (int i = 0; i < size_residual; ++i) {
      for (int j = 0; j < size_rt; ++j) {
        ASSERT_NEAR(residual_adiff[i].derivatives()(size_point + j),
                    jac_rt[i * size_rt + j], 1e-15);
      }
    }
  }

  template <int N>
  void RunTest(const ProjectionType& type, const double* camera) {
    constexpr int size = N;

    // Autodiff-ed version will be used as reference/expected values
    AScalar camera_adiff[size];
    SetupADiff(size, &camera[0], &camera_adiff[0]);
    ReprojectionError2D autodiff(type, observed, scale);
    autodiff(camera_adiff, rt_adiff, point_adiff, residual_adiff);

    // We test for analytic evaluation
    double jac_camera[size_residual * size];
    const double* params[] = {camera, rt, point};
    double* jacobians[] = {jac_camera, jac_rt, jac_point};
    ReprojectionError2DAnalytic<size> analytic(type, observed, scale);
    analytic.Evaluate(params, residuals, &jacobians[0]);

    // Check
    CheckJacobians(size, jac_camera);
  }

  constexpr static int size_residual = 2;
  constexpr static int size_point = 3;
  constexpr static int size_rt = 6;

  Vec2d observed;
  double scale{0.1};
  const double point[size_point] = {1.0, 2.0, 3.0};
  const double rt[size_rt] = {0.1, 0.2, 0.3, 0.4, 0.5, 0.6};

  AScalar residual_adiff[size_residual];
  AScalar point_adiff[size_point];
  AScalar rt_adiff[size_rt];

  double residuals[size_residual];
  double jac_rt[size_residual * size_rt];
  double jac_point[size_residual * size_point];
};

TEST_F(ReprojectionError2DFixture, BrownAnalyticErrorEvaluatesOK) {
  constexpr int size = 9;

  // focal, ar, cx, cy, k1, k2, k3, p1, p2
  const double camera[size] = {0.3, 1.0, 0.001, -0.02, 0.1, -0.03, 0.001, -0.005, 0.001};
  RunTest<size>(ProjectionType::BROWN, &camera[0]);
}

TEST_F(ReprojectionError2DFixture, PerspectiveAnalyticErrorEvaluatesOK) {
  constexpr int size = 3;

  // focal k1, k2
  const double camera[size] = {0.3, 0.1, -0.03};
  RunTest<size>(ProjectionType::PERSPECTIVE, &camera[0]);
}

TEST_F(ReprojectionError2DFixture, FisheyeAnalyticErrorEvaluatesOK) {
  constexpr int size = 3;

  // focal k1, k2
  const double camera[size] = {0.3, 0.1, -0.03};
  RunTest<size>(ProjectionType::FISHEYE, &camera[0]);
}

TEST_F(ReprojectionError2DFixture, DualAnalyticErrorEvaluatesOK) {
  constexpr int size = 4;

  // focal k1, k2
  const double camera[size] = {0.5, 0.3, 0.1, -0.03};
  RunTest<size>(ProjectionType::DUAL, &camera[0]);
}

class ReprojectionError3DFixture : public ::testing::Test {
 public:
 static constexpr int size = 3;

 typedef Eigen::AutoDiffScalar<Eigen::VectorXd> AScalar;
 ReprojectionError3DFixture() { observed << 0.5, 0.5; }

 void SetupADiff() {
   const int total_size = size_point + size_rt;
   for (int i = 0; i < size_point; ++i) {
     point_adiff[i].value() = point[i];
     point_adiff[i].derivatives() = VecXd::Unit(total_size, i);
   }
   for (int i = 0; i < size_rt; ++i) {
     rt_adiff[i].value() = rt[i];
     rt_adiff[i].derivatives() = VecXd::Unit(total_size, size_point + i);
   }
  }

  void CheckJacobians() {
    for (int i = 0; i < size; ++i) {
      for (int j = 0; j < size_rt; ++j) {
        ASSERT_NEAR(residual_adiff[i].derivatives()(size_point + j),
                    jac_rt[i * size_rt + j], 1e-14);
      }
    }
  }

  constexpr static int size_point = 3;
  constexpr static int size_rt = 6;

  Vec2d observed;
  double scale{0.1};
  const double point[size_point] = {1.0, 2.0, 3.0};
  const double rt[size_rt] = {0.1, 0.2, 0.3, 0.4, 0.5, 0.6};

  AScalar residual_adiff[size];
  AScalar point_adiff[size_point];
  AScalar rt_adiff[size_rt];

  double residuals[size];
  double jac_rt[size * size_rt];
  double jac_point[size * size_point];
};

TEST_F(ReprojectionError3DFixture, AnalyticErrorEvaluatesOK) {
  // Autodiff-ed version will be used as reference/expected values
    SetupADiff();
    AScalar dummy_adiff;
    ReprojectionError3D autodiff(ProjectionType::SPHERICAL, observed, scale);
    autodiff(&dummy_adiff, rt_adiff, point_adiff, residual_adiff);

    // We test for analytic evaluation
    double dummy = 0.;
    double dummy_jac[] = {0., 0., 0.};
    const double* params[] = {&dummy, rt, point};
    double* jacobians[] = {&dummy_jac[0], jac_rt, jac_point};
    ReprojectionError3DAnalytic analytic(ProjectionType::SPHERICAL, observed, scale);
    analytic.Evaluate(params, residuals, &jacobians[0]);

    // Check
    CheckJacobians();
}

<commit_msg>chore: doc<commit_after>#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include <Eigen/Dense>
#include <unsupported/Eigen/AutoDiff>

#include <bundle/src/projection_errors.h>

class ReprojectionError2DFixture : public ::testing::Test {
 public:
  typedef Eigen::AutoDiffScalar<Eigen::VectorXd> AScalar;
  ReprojectionError2DFixture() { observed << 0.5, 0.5; }

  void SetupADiff(int size, const double* camera, AScalar* camera_adiff) {
    const int total_size = size_point + size_rt + size;
    for (int i = 0; i < size_point; ++i) {
      point_adiff[i].value() = point[i];
      point_adiff[i].derivatives() = VecXd::Unit(total_size, i);
    }
    for (int i = 0; i < size_rt; ++i) {
      rt_adiff[i].value() = rt[i];
      rt_adiff[i].derivatives() = VecXd::Unit(total_size, size_point + i);
    }
    for (int i = 0; i < size; ++i) {
      camera_adiff[i].value() = camera[i];
      camera_adiff[i].derivatives() = VecXd::Unit(total_size, size_point + size_rt + i);
    }
  }

  void CheckJacobians(int size, const double* jac_camera) {
    for (int i = 0; i < size_residual; ++i) {
      for (int j = 0; j < size; ++j) {
        ASSERT_NEAR(residual_adiff[i].derivatives()(size_point + size_rt + j),
                    jac_camera[i * size + j], 1e-15);
      }
    }
    for (int i = 0; i < size_residual; ++i) {
      for (int j = 0; j < size_point; ++j) {
        ASSERT_NEAR(residual_adiff[i].derivatives()(j),
                    jac_point[i * size_point + j], 1e-15);
      }
    }
    for (int i = 0; i < size_residual; ++i) {
      for (int j = 0; j < size_rt; ++j) {
        ASSERT_NEAR(residual_adiff[i].derivatives()(size_point + j),
                    jac_rt[i * size_rt + j], 1e-15);
      }
    }
  }

  template <int N>
  void RunTest(const ProjectionType& type, const double* camera) {
    constexpr int size = N;

    // Autodiff-ed version will be used as reference/expected values
    AScalar camera_adiff[size];
    SetupADiff(size, &camera[0], &camera_adiff[0]);
    ReprojectionError2D autodiff(type, observed, scale);
    autodiff(camera_adiff, rt_adiff, point_adiff, residual_adiff);

    // We test for analytic evaluation
    double jac_camera[size_residual * size];
    const double* params[] = {camera, rt, point};
    double* jacobians[] = {jac_camera, jac_rt, jac_point};
    ReprojectionError2DAnalytic<size> analytic(type, observed, scale);
    analytic.Evaluate(params, residuals, &jacobians[0]);

    // Check
    CheckJacobians(size, jac_camera);
  }

  constexpr static int size_residual = 2;
  constexpr static int size_point = 3;
  constexpr static int size_rt = 6;

  Vec2d observed;
  double scale{0.1};
  const double point[size_point] = {1.0, 2.0, 3.0};
  const double rt[size_rt] = {0.1, 0.2, 0.3, 0.4, 0.5, 0.6};

  AScalar residual_adiff[size_residual];
  AScalar point_adiff[size_point];
  AScalar rt_adiff[size_rt];

  double residuals[size_residual];
  double jac_rt[size_residual * size_rt];
  double jac_point[size_residual * size_point];
};

TEST_F(ReprojectionError2DFixture, BrownAnalyticErrorEvaluatesOK) {
  constexpr int size = 9;

  // focal, ar, cx, cy, k1, k2, k3, p1, p2
  const double camera[size] = {0.3, 1.0, 0.001, -0.02, 0.1, -0.03, 0.001, -0.005, 0.001};
  RunTest<size>(ProjectionType::BROWN, &camera[0]);
}

TEST_F(ReprojectionError2DFixture, PerspectiveAnalyticErrorEvaluatesOK) {
  constexpr int size = 3;

  // focal k1, k2
  const double camera[size] = {0.3, 0.1, -0.03};
  RunTest<size>(ProjectionType::PERSPECTIVE, &camera[0]);
}

TEST_F(ReprojectionError2DFixture, FisheyeAnalyticErrorEvaluatesOK) {
  constexpr int size = 3;

  // focal k1, k2
  const double camera[size] = {0.3, 0.1, -0.03};
  RunTest<size>(ProjectionType::FISHEYE, &camera[0]);
}

TEST_F(ReprojectionError2DFixture, DualAnalyticErrorEvaluatesOK) {
  constexpr int size = 4;

  // transtion, focal, k1, k2
  const double camera[size] = {0.5, 0.3, 0.1, -0.03};
  RunTest<size>(ProjectionType::DUAL, &camera[0]);
}

class ReprojectionError3DFixture : public ::testing::Test {
 public:
 static constexpr int size = 3;

 typedef Eigen::AutoDiffScalar<Eigen::VectorXd> AScalar;
 ReprojectionError3DFixture() { observed << 0.5, 0.5; }

 void SetupADiff() {
   const int total_size = size_point + size_rt;
   for (int i = 0; i < size_point; ++i) {
     point_adiff[i].value() = point[i];
     point_adiff[i].derivatives() = VecXd::Unit(total_size, i);
   }
   for (int i = 0; i < size_rt; ++i) {
     rt_adiff[i].value() = rt[i];
     rt_adiff[i].derivatives() = VecXd::Unit(total_size, size_point + i);
   }
  }

  void CheckJacobians() {
    for (int i = 0; i < size; ++i) {
      for (int j = 0; j < size_rt; ++j) {
        ASSERT_NEAR(residual_adiff[i].derivatives()(size_point + j),
                    jac_rt[i * size_rt + j], 1e-14);
      }
    }
  }

  constexpr static int size_point = 3;
  constexpr static int size_rt = 6;

  Vec2d observed;
  double scale{0.1};
  const double point[size_point] = {1.0, 2.0, 3.0};
  const double rt[size_rt] = {0.1, 0.2, 0.3, 0.4, 0.5, 0.6};

  AScalar residual_adiff[size];
  AScalar point_adiff[size_point];
  AScalar rt_adiff[size_rt];

  double residuals[size];
  double jac_rt[size * size_rt];
  double jac_point[size * size_point];
};

TEST_F(ReprojectionError3DFixture, AnalyticErrorEvaluatesOK) {
  // Autodiff-ed version will be used as reference/expected values
    SetupADiff();
    AScalar dummy_adiff;
    ReprojectionError3D autodiff(ProjectionType::SPHERICAL, observed, scale);
    autodiff(&dummy_adiff, rt_adiff, point_adiff, residual_adiff);

    // We test for analytic evaluation
    double dummy = 0.;
    double dummy_jac[] = {0., 0., 0.};
    const double* params[] = {&dummy, rt, point};
    double* jacobians[] = {&dummy_jac[0], jac_rt, jac_point};
    ReprojectionError3DAnalytic analytic(ProjectionType::SPHERICAL, observed, scale);
    analytic.Evaluate(params, residuals, &jacobians[0]);

    // Check
    CheckJacobians();
}

<|endoftext|>


<commit_before>/**
 *  Copyright (c) 2014-present, Facebook, Inc.
 *  All rights reserved.
 *
 *  This source code is licensed under both the Apache 2.0 license (found in the
 *  LICENSE file in the root directory of this source tree) and the GPLv2 (found
 *  in the COPYING file in the root directory of this source tree).
 *  You may select, at your option, one of the above-listed licenses.
 */

#include <cmath>

#include <boost/property_tree/ptree.hpp>
#include <boost/property_tree/xml_parser.hpp>

#include <osquery/config.h>
#include <osquery/core.h>
#include <osquery/database.h>
#include <osquery/logger.h>
#include <osquery/registry_factory.h>
#include <osquery/tables.h>

#include "osquery/config/parsers/feature_vectors.h"
#include "osquery/core/conversions.h"
#include "osquery/core/json.h"
#include "osquery/core/windows/wmi.h"
#include "osquery/events/windows/windows_event_log.h"
#include "osquery/filesystem/fileops.h"

namespace pt = boost::property_tree;

namespace osquery {

const std::string kPowershellDomain{"powershell_script_block_logs"};
const std::string kScriptBlockPrefix{"script_block."};
const std::string kScriptBlocksReceivedSuffix{".blocks_received"};
const std::wstring kPowershellEventsChannel{
    L"microsoft-windows-powershell/operational"};

const int kScriptBlockLoggingEid{4104};
const int kCharFreqVectorLen{255};

void parseTree(const pt::ptree& tree, std::map<std::string, std::string>& res);

class PowershellEventSubscriber
    : public EventSubscriber<WindowsEventLogEventPublisher> {
 public:
  Status init() override {
    // Before starting our subscription, purge any residual db entries as it's
    // unlikely we'll finish re-assmebling them
    std::vector<std::string> keys;
    scanDatabaseKeys(kEvents, keys, kScriptBlockPrefix);
    for (const auto& k : keys) {
      auto s = deleteDatabaseValue(kEvents, k);
      if (!s.ok()) {
        VLOG(1) << "Failed to delete stale script block from the database "
                << k;
      }
    }

    auto wc = createSubscriptionContext();
    wc->sources.insert(kPowershellEventsChannel);

    subscribe(&PowershellEventSubscriber::Callback, wc);
    return Status();
  }

  Status Callback(const ECRef& ec, const SCRef& sc);

  void addScriptResult(Row& results);
};

REGISTER(PowershellEventSubscriber, "event_subscriber", "powershell_events");

inline double cosineSimilarity(std::vector<double>& script_freqs,
                               std::vector<double>& global_freqs) {
  auto dot = 0.0;
  auto mag1 = 0.0;
  auto mag2 = 0.0;

  for (size_t i = 0; i < global_freqs.size(); i++) {
    dot += script_freqs[i] * global_freqs[i];
    mag1 += script_freqs[i] * script_freqs[i];
    mag2 += global_freqs[i] * global_freqs[i];
  }

  mag1 = sqrt(mag1);
  mag2 = sqrt(mag2);

  return dot / (mag1 * mag2);
}

void PowershellEventSubscriber::addScriptResult(Row& results) {
  Row r;
  r["time"] = results["time"];
  r["datetime"] = results["datetime"];
  r["script_block_id"] = results["ScriptBlockId"];
  r["script_block_count"] = INTEGER(results["MessageTotal"]);
  r["script_name"] = results["Name"];
  r["script_path"] = results["Path"];
  r["script_text"] = results["ScriptBlockText"];

  // Grab the feature vectors from the configuration
  auto parser = Config::getParser(kFeatureVectorsRootKey);
  if (parser == nullptr) {
    VLOG(1) << "Failed to get configured feature vectors";
    add(r);
    return;
  }

  // Get the reassembled powershell scripts character frequency vector
  std::vector<double> freqs(kCharFreqVectorLen, 0.0);
  for (const auto chr : r["script_text"]) {
    if (chr < kCharFreqVectorLen) {
      freqs[chr] += 1.0 / r["script_text"].length();
    }
  }

  const auto& cf =
      parser->getData().doc()["feature_vectors"]["character_frequencies"];
  if (cf.Empty() || cf.Size() != kCharFreqVectorLen) {
    VLOG(1) << "Invalid character frequency map found, using default values";
    add(r);
    return;
  }

  std::vector<double> cfg_freqs(kCharFreqVectorLen, 0.0);
  for (unsigned int i = 0; i < cf.Size(); i++) {
    cfg_freqs[i] = cf[i].GetDouble();
  }
  r["cosine_similarity"] = DOUBLE(cosineSimilarity(freqs, cfg_freqs));
  add(r);
}

Status PowershellEventSubscriber::Callback(const ECRef& ec, const SCRef& sc) {
  // For script block logging we only care about events with script blocks
  auto eid = ec->eventRecord.get("Event.System.EventID", -1);
  if (eid != kScriptBlockLoggingEid) {
    return Status();
  }

  Row results;
  for (const auto& node : ec->eventRecord.get_child("Event", pt::ptree())) {
    if (node.first == "System" || node.first == "<xmlattr>") {
      continue;
    }
    // #4357: This should make use of RapidJSON
    parseTree(node.second, results);
  }

  FILETIME etime;
  GetSystemTimeAsFileTime(&etime);
  results["time"] = BIGINT(filetimeToUnixtime(etime));
  results["datetime"] =
      ec->eventRecord.get("Event.System.TimeCreated.<xmlattr>.SystemTime", "");

  // If there's only one script block no reassembly is needed
  if (results["MessageTotal"] == "1") {
    addScriptResult(results);
    return Status();
  }

  // Add the script content to the DB for later reassembly
  auto s = setDatabaseValue(kEvents,
                            kScriptBlockPrefix + results["ScriptBlockId"] +
                                "." + results["MessageNumber"],
                            results["ScriptBlockText"]);
  if (!s.ok()) {
    LOG(WARNING) << "Failed to add new Powershell block to database for script "
                 << results["ScriptBlockId"];
  }

  // If we expect more blocks bail out early
  if (results["MessageNumber"] != results["MessageTotal"]) {
    return Status();
  }

  // Otherwise all script blocks should be accounted for so reconstruct
  std::vector<std::string> keys;
  s = scanDatabaseKeys(
      kEvents, keys, kScriptBlockPrefix + results["ScriptBlockId"]);
  if (!s.ok()) {
    LOG(WARNING) << "Failed to look up powershell script blocks for "
                 << results["ScriptBlockId"];
    return Status(1);
  }

  std::string powershell_script{""};
  for (const auto& key : keys) {
    std::string val{""};
    s = getDatabaseValue(kEvents, key, val);
    if (!s.ok()) {
      LOG(WARNING) << "Failed to retrieve script block " << key;
      continue;
    }

    powershell_script += val;

    s = deleteDatabaseValue(kEvents, key);
    if (!s.ok()) {
      LOG(WARNING) << "Failed to delete script block key from db " << key;
    }
  }

  results["ScriptBlockText"] = powershell_script;
  addScriptResult(results);

  return Status();
}
} // namespace osquery
<commit_msg>events: updating logging verbiage and off-by-one for powershell events (#4502)<commit_after>/**
 *  Copyright (c) 2014-present, Facebook, Inc.
 *  All rights reserved.
 *
 *  This source code is licensed under both the Apache 2.0 license (found in the
 *  LICENSE file in the root directory of this source tree) and the GPLv2 (found
 *  in the COPYING file in the root directory of this source tree).
 *  You may select, at your option, one of the above-listed licenses.
 */

#include <cmath>

#include <boost/property_tree/ptree.hpp>
#include <boost/property_tree/xml_parser.hpp>

#include <osquery/config.h>
#include <osquery/core.h>
#include <osquery/database.h>
#include <osquery/logger.h>
#include <osquery/registry_factory.h>
#include <osquery/tables.h>

#include "osquery/config/parsers/feature_vectors.h"
#include "osquery/core/conversions.h"
#include "osquery/core/json.h"
#include "osquery/core/windows/wmi.h"
#include "osquery/events/windows/windows_event_log.h"
#include "osquery/filesystem/fileops.h"

namespace pt = boost::property_tree;

namespace osquery {

const std::string kPowershellDomain{"powershell_script_block_logs"};
const std::string kScriptBlockPrefix{"script_block."};
const std::string kScriptBlocksReceivedSuffix{".blocks_received"};
const std::wstring kPowershellEventsChannel{
    L"microsoft-windows-powershell/operational"};

const int kScriptBlockLoggingEid{4104};
const int kCharFreqVectorLen{256};

void parseTree(const pt::ptree& tree, std::map<std::string, std::string>& res);

class PowershellEventSubscriber
    : public EventSubscriber<WindowsEventLogEventPublisher> {
 public:
  Status init() override {
    // Before starting our subscription, purge any residual db entries as it's
    // unlikely we'll finish re-assmebling them
    std::vector<std::string> keys;
    scanDatabaseKeys(kEvents, keys, kScriptBlockPrefix);
    for (const auto& k : keys) {
      auto s = deleteDatabaseValue(kEvents, k);
      if (!s.ok()) {
        VLOG(1) << "Failed to delete stale script block from the database "
                << k;
      }
    }

    auto wc = createSubscriptionContext();
    wc->sources.insert(kPowershellEventsChannel);

    subscribe(&PowershellEventSubscriber::Callback, wc);
    return Status();
  }

  Status Callback(const ECRef& ec, const SCRef& sc);

  void addScriptResult(Row& results);
};

REGISTER(PowershellEventSubscriber, "event_subscriber", "powershell_events");

inline double cosineSimilarity(std::vector<double>& script_freqs,
                               std::vector<double>& global_freqs) {
  auto dot = 0.0;
  auto mag1 = 0.0;
  auto mag2 = 0.0;

  for (size_t i = 0; i < global_freqs.size(); i++) {
    dot += script_freqs[i] * global_freqs[i];
    mag1 += script_freqs[i] * script_freqs[i];
    mag2 += global_freqs[i] * global_freqs[i];
  }

  mag1 = sqrt(mag1);
  mag2 = sqrt(mag2);

  return dot / (mag1 * mag2);
}

void PowershellEventSubscriber::addScriptResult(Row& results) {
  Row r;
  r["time"] = results["time"];
  r["datetime"] = results["datetime"];
  r["script_block_id"] = results["ScriptBlockId"];
  r["script_block_count"] = INTEGER(results["MessageTotal"]);
  r["script_name"] = results["Name"];
  r["script_path"] = results["Path"];
  r["script_text"] = results["ScriptBlockText"];

  // Grab the feature vectors from the configuration
  auto parser = Config::getParser(kFeatureVectorsRootKey);
  if (parser == nullptr) {
    VLOG(1) << "Failed to get configured feature vectors";
    add(r);
    return;
  }

  // Get the reassembled powershell scripts character frequency vector
  std::vector<double> freqs(kCharFreqVectorLen, 0.0);
  for (const auto chr : r["script_text"]) {
    if (chr < kCharFreqVectorLen) {
      freqs[chr] += 1.0 / r["script_text"].length();
    }
  }

  const auto& cf =
      parser->getData().doc()["feature_vectors"]["character_frequencies"];
  if (cf.Empty() || cf.Size() != kCharFreqVectorLen) {
    VLOG(1) << "Invalid character frequency map found, skipping computation";
    add(r);
    return;
  }

  std::vector<double> cfg_freqs(kCharFreqVectorLen, 0.0);
  for (unsigned int i = 0; i < cf.Size(); i++) {
    cfg_freqs[i] = cf[i].GetDouble();
  }
  r["cosine_similarity"] = DOUBLE(cosineSimilarity(freqs, cfg_freqs));
  add(r);
}

Status PowershellEventSubscriber::Callback(const ECRef& ec, const SCRef& sc) {
  // For script block logging we only care about events with script blocks
  auto eid = ec->eventRecord.get("Event.System.EventID", -1);
  if (eid != kScriptBlockLoggingEid) {
    return Status();
  }

  Row results;
  for (const auto& node : ec->eventRecord.get_child("Event", pt::ptree())) {
    if (node.first == "System" || node.first == "<xmlattr>") {
      continue;
    }
    // #4357: This should make use of RapidJSON
    parseTree(node.second, results);
  }

  FILETIME etime;
  GetSystemTimeAsFileTime(&etime);
  results["time"] = BIGINT(filetimeToUnixtime(etime));
  results["datetime"] =
      ec->eventRecord.get("Event.System.TimeCreated.<xmlattr>.SystemTime", "");

  // If there's only one script block no reassembly is needed
  if (results["MessageTotal"] == "1") {
    addScriptResult(results);
    return Status();
  }

  // Add the script content to the DB for later reassembly
  auto s = setDatabaseValue(kEvents,
                            kScriptBlockPrefix + results["ScriptBlockId"] +
                                "." + results["MessageNumber"],
                            results["ScriptBlockText"]);
  if (!s.ok()) {
    LOG(WARNING) << "Failed to add new Powershell block to database for script "
                 << results["ScriptBlockId"];
  }

  // If we expect more blocks bail out early
  if (results["MessageNumber"] != results["MessageTotal"]) {
    return Status();
  }

  // Otherwise all script blocks should be accounted for so reconstruct
  std::vector<std::string> keys;
  s = scanDatabaseKeys(
      kEvents, keys, kScriptBlockPrefix + results["ScriptBlockId"]);
  if (!s.ok()) {
    LOG(WARNING) << "Failed to look up powershell script blocks for "
                 << results["ScriptBlockId"];
    return Status(1);
  }

  std::string powershell_script{""};
  for (const auto& key : keys) {
    std::string val{""};
    s = getDatabaseValue(kEvents, key, val);
    if (!s.ok()) {
      LOG(WARNING) << "Failed to retrieve script block " << key;
      continue;
    }

    powershell_script += val;

    s = deleteDatabaseValue(kEvents, key);
    if (!s.ok()) {
      LOG(WARNING) << "Failed to delete script block key from db " << key;
    }
  }

  results["ScriptBlockText"] = powershell_script;
  addScriptResult(results);

  return Status();
}
} // namespace osquery
<|endoftext|>


<commit_before>// -*- C++ -*-
//===------------------------- fuzz_test.cpp ------------------------------===//
//
//                     The LLVM Compiler Infrastructure
//
// This file is dual licensed under the MIT and the University of Illinois Open
// Source Licenses. See LICENSE.TXT for details.
//
//===----------------------------------------------------------------------===//

//  A simple program for running regressions on the fuzzing routines.
//  This code is not part of any shipping product.
//
//  To build:
//      clang++ -std=c++11 fuzz_test.cpp fuzzing.cpp
//
//  To use:
//      fuzz_test -r partial_sort [-v] files...
//
//  Each file should contain a test case.

#include <iostream>
#include <fstream>
#include <vector>
#include <map>
#include <chrono>

#include "fuzzing.h"

typedef int (*FuzzProc) (const uint8_t *data, size_t size);

const std::map<std::string, FuzzProc> procs = {
    {"sort",                fuzzing::sort},
    {"stable_sort",         fuzzing::stable_sort},
    {"partition",           fuzzing::partition},
    {"partition_copy",      fuzzing::partition_copy},
    {"stable_partition",    fuzzing::stable_partition},
    {"unique",              fuzzing::unique},
    {"unique_copy",         fuzzing::unique_copy},
    {"nth_element",         fuzzing::nth_element},
    {"partial_sort",        fuzzing::partial_sort},
    {"partial_sort_copy",   fuzzing::partial_sort_copy},
    {"make_heap",           fuzzing::make_heap},
    {"push_heap",           fuzzing::push_heap},
    {"pop_heap",            fuzzing::pop_heap},
    {"regex_ECMAScript",    fuzzing::regex_ECMAScript},
    {"regex_POSIX",         fuzzing::regex_POSIX},
    {"regex_extended",      fuzzing::regex_extended},
    {"regex_awk",           fuzzing::regex_awk},
    {"regex_grep",          fuzzing::regex_grep},
    {"regex_egrep",         fuzzing::regex_egrep},
    {"search",              fuzzing::search}
};

bool verbose = false;

void test_one(const char *filename, FuzzProc fp)
{
    std::vector<uint8_t> v;
    std::ifstream f (filename, std::ios::binary);
    if (!f.is_open())
        std::cerr << "## Can't open '" << filename << "'" << std::endl;
    else {
        typedef std::istream_iterator<uint8_t> Iter;
        std::copy(Iter(f), Iter(), std::back_inserter(v));
        if (verbose)
            std::cout << "File '" << filename << "' contains " << v.size() << " entries" << std::endl;
        const auto start_time = std::chrono::steady_clock::now();
        int ret = fp (v.data(), v.size());
        const auto finish_time = std::chrono::steady_clock::now();
        if (ret != 0)
            std::cerr << "## Failure code: " << ret << std::endl;
        if (verbose)
            std::cout << "Execution time: "
                << std::chrono::duration_cast<std::chrono::milliseconds>(finish_time - start_time).count()
                << " milliseconds" << std::endl;
        }
}

void usage (const char *name)
{
    std::cout << "Usage: " << name << " -r proc [-v] files..." << std::endl;
    std::cout << "Supported routines:" << std::endl;
    for (const auto &p : procs)
        std::cout << "    " << p.first << std::endl;
    std::cout << std::endl;
}

// Poor man's command-line options
const std::string dashR("-r");
const std::string dashV("-v");

int main(int argc, char *argv[])
{
    if (argc < 4 || dashR != argv[1] || procs.find(argv[2]) == procs.end())
        usage(argv[0]);
    else {
        FuzzProc fp = procs.find(argv[2])->second;
        int firstFile = 3;
        if (dashV == argv[firstFile])
        {
            verbose = true;
            ++firstFile;
        }
        for (int i = firstFile; i < argc; ++i)
            test_one(argv[i], fp);
        }
}
<commit_msg>Use high_resolution_clock instead of steady_clock. Also now builds with gcc 7.2 (for comparison purposes)<commit_after>// -*- C++ -*-
//===------------------------- fuzz_test.cpp ------------------------------===//
//
//                     The LLVM Compiler Infrastructure
//
// This file is dual licensed under the MIT and the University of Illinois Open
// Source Licenses. See LICENSE.TXT for details.
//
//===----------------------------------------------------------------------===//

//  A simple program for running regressions on the fuzzing routines.
//  This code is not part of any shipping product.
//
//  To build:
//      clang++ -std=c++11 fuzz_test.cpp fuzzing.cpp
//
//  To use:
//      fuzz_test -r partial_sort [-v] files...
//
//  Each file should contain a test case.

#include <iostream>
#include <fstream>
#include <iterator>
#include <vector>
#include <map>
#include <chrono>

#include "fuzzing.h"

typedef int (*FuzzProc) (const uint8_t *data, size_t size);

const std::map<std::string, FuzzProc> procs = {
    {"sort",                fuzzing::sort},
    {"stable_sort",         fuzzing::stable_sort},
    {"partition",           fuzzing::partition},
    {"partition_copy",      fuzzing::partition_copy},
    {"stable_partition",    fuzzing::stable_partition},
    {"unique",              fuzzing::unique},
    {"unique_copy",         fuzzing::unique_copy},
    {"nth_element",         fuzzing::nth_element},
    {"partial_sort",        fuzzing::partial_sort},
    {"partial_sort_copy",   fuzzing::partial_sort_copy},
    {"make_heap",           fuzzing::make_heap},
    {"push_heap",           fuzzing::push_heap},
    {"pop_heap",            fuzzing::pop_heap},
    {"regex_ECMAScript",    fuzzing::regex_ECMAScript},
    {"regex_POSIX",         fuzzing::regex_POSIX},
    {"regex_extended",      fuzzing::regex_extended},
    {"regex_awk",           fuzzing::regex_awk},
    {"regex_grep",          fuzzing::regex_grep},
    {"regex_egrep",         fuzzing::regex_egrep},
    {"search",              fuzzing::search}
};

bool verbose = false;

void test_one(const char *filename, FuzzProc fp)
{
    std::vector<uint8_t> v;
    std::ifstream f (filename, std::ios::binary);
    if (!f.is_open())
        std::cerr << "## Can't open '" << filename << "'" << std::endl;
    else {
        typedef std::istream_iterator<uint8_t> Iter;
        std::copy(Iter(f), Iter(), std::back_inserter(v));
        if (verbose)
            std::cout << "File '" << filename << "' contains " << v.size() << " entries" << std::endl;
        const auto start_time = std::chrono::high_resolution_clock::now();
        int ret = fp (v.data(), v.size());
        const auto finish_time = std::chrono::high_resolution_clock::now();
        if (ret != 0)
            std::cerr << "## Failure code: " << ret << std::endl;
        if (verbose)
            std::cout << "Execution time: "
                << std::chrono::duration_cast<std::chrono::milliseconds>(finish_time - start_time).count()
                << " milliseconds" << std::endl;
        }
}

void usage (const char *name)
{
    std::cout << "Usage: " << name << " -r proc [-v] files..." << std::endl;
    std::cout << "Supported routines:" << std::endl;
    for (const auto &p : procs)
        std::cout << "    " << p.first << std::endl;
    std::cout << std::endl;
}

// Poor man's command-line options
const std::string dashR("-r");
const std::string dashV("-v");

int main(int argc, char *argv[])
{
    if (argc < 4 || dashR != argv[1] || procs.find(argv[2]) == procs.end())
        usage(argv[0]);
    else {
        FuzzProc fp = procs.find(argv[2])->second;
        int firstFile = 3;
        if (dashV == argv[firstFile])
        {
            verbose = true;
            ++firstFile;
        }
        for (int i = firstFile; i < argc; ++i)
            test_one(argv[i], fp);
        }
}
<|endoftext|>