<commit_before>// ***************************************************************************
// FastaIndex.cpp (c) 2010 Erik Garrison <erik.garrison@bc.edu>
// Marth Lab, Department of Biology, Boston College
// All rights reserved.
// ---------------------------------------------------------------------------
// Last modified: 9 February 2010 (EG)
// ---------------------------------------------------------------------------

#include "Fasta.h"

FastaIndexEntry::FastaIndexEntry(string name, int length, long long offset, int line_blen, int line_len)
    : name(name)
    , length(length)
    , offset(offset)
    , line_blen(line_blen)
    , line_len(line_len)
{}

FastaIndexEntry::FastaIndexEntry(void) // empty constructor
{ clear(); }

FastaIndexEntry::~FastaIndexEntry(void)
{}

void FastaIndexEntry::clear(void)
{
    name = "";
    length = NULL;
    offset = -1;  // no real offset will ever be below 0, so this allows us to
                  // check if we have already recorded a real offset
    line_blen = NULL;
    line_len = NULL;
}

ostream& operator<<(ostream& output, const FastaIndexEntry& e) {
    // just write the first component of the name, for compliance with other tools
    output << split(e.name, ' ').at(0) << "\t" << e.length << "\t" << e.offset << "\t" <<
        e.line_blen << "\t" << e.line_len;
    return output;  // for multiple << operators.
}

FastaIndex::FastaIndex(void)
{}

void FastaIndex::readIndexFile(string fname) {
    string line;
    long long linenum = 0;
    indexFile.open(fname.c_str(), ifstream::in);
    if (indexFile.is_open()) {
        while (getline (indexFile, line)) {
            ++linenum;
            // the fai format defined in samtools is tab-delimited, every line being:
            // fai->name[i], (int)x.len, (long long)x.offset, (int)x.line_blen, (int)x.line_len
            vector<string> fields = split(line, '\t');
            if (fields.size() == 5) {  // if we don't get enough fields then there is a problem with the file
                // note that fields[0] is the sequence name
                char* end;
                string name = split(fields[0], " \t").at(0);  // key by first token of name
                sequenceNames.push_back(name);
                this->insert(make_pair(name, FastaIndexEntry(fields[0], atoi(fields[1].c_str()),
                                                    strtoll(fields[2].c_str(), &end, 10),
                                                    atoi(fields[3].c_str()),
                                                    atoi(fields[4].c_str()))));
            } else {
                cerr << "Warning: malformed fasta index file " << fname <<
                    "does not have enough fields @ line " << linenum << endl;
                cerr << line << endl;
                exit(1);
            }
        }
    } else {
        cerr << "could not open index file " << fname << endl;
        exit(1);
    }
}

// for consistency this should be a class method
bool fastaIndexEntryCompare ( FastaIndexEntry a, FastaIndexEntry b) { return (a.offset<b.offset); }

ostream& operator<<(ostream& output, FastaIndex& fastaIndex) {
    vector<FastaIndexEntry> sortedIndex;
    for(vector<string>::const_iterator it = fastaIndex.sequenceNames.begin(); it != fastaIndex.sequenceNames.end(); ++it)
    {
        sortedIndex.push_back(fastaIndex[*it]);
    }
    sort(sortedIndex.begin(), sortedIndex.end(), fastaIndexEntryCompare);
    for( vector<FastaIndexEntry>::iterator fit = sortedIndex.begin(); fit != sortedIndex.end(); ++fit) {
        output << *fit << endl;
    }
    return output;
}

void FastaIndex::indexReference(string refname) {
    // overview:
    //  for line in the reference fasta file
    //  track byte offset from the start of the file
    //  if line is a fasta header, take the name and dump the last sequnece to the index
    //  if line is a sequence, add it to the current sequence
    //cerr << "indexing fasta reference " << refname << endl;
    string line;
    FastaIndexEntry entry;  // an entry buffer used in processing
    entry.clear();
    int line_length = 0;
    long long offset = 0;  // byte offset from start of file
    long long line_number = 0; // current line number
    bool mismatchedLineLengths = false; // flag to indicate if our line length changes mid-file
                                        // this will be used to raise an error
                                        // if we have a line length change at
                                        // any line other than the last line in
                                        // the sequence
    bool emptyLine = false;  // flag to catch empty lines, which we allow for
                             // index generation only on the last line of the sequence
    ifstream refFile;
    refFile.open(refname.c_str());
    if (refFile.is_open()) {
        while (getline(refFile, line)) {
            ++line_number;
            line_length = line.length();
            if (line[0] == ';') {
                // fasta comment, skip
            } else if (line[0] == '+') {
                // fastq quality header
                getline(refFile, line);
                line_length = line.length();
                offset += line_length + 1;
                // get and don't handle the quality line
                getline(refFile, line);
                line_length = line.length();
            } else if (line[0] == '>' || line[0] == '@') { // fasta /fastq header
                // if we aren't on the first entry, push the last sequence into the index
                if (entry.name != "") {
                    mismatchedLineLengths = false; // reset line length error tracker for every new sequence
                    emptyLine = false;
                    flushEntryToIndex(entry);
                    entry.clear();
                }
                entry.name = line.substr(1, line_length - 1);
            } else { // we assume we have found a sequence line
                if (entry.offset == -1) // NB initially the offset is -1
                    entry.offset = offset;
                entry.length += line_length;
                if (entry.line_len) {
                    //entry.line_len = entry.line_len ? entry.line_len : line_length + 1;
                    if (mismatchedLineLengths || emptyLine) {
                        if (line_length == 0) {
                            emptyLine = true; // flag empty lines, raise error only if this is embedded in the sequence
                        } else {
                            if (emptyLine) {
                                cerr << "ERROR: embedded newline";
                            } else {
                                cerr << "ERROR: mismatched line lengths";
                            }
                            cerr << " at line " << line_number << " within sequence " << entry.name <<
                                endl << "File not suitable for fasta index generation." << endl;
                            exit(1);
                        }
                    }
                    // this flag is set here and checked on the next line
                    // because we may have reached the end of the sequence, in
                    // which case a mismatched line length is OK
                    if (entry.line_len != line_length + 1) {
                        mismatchedLineLengths = true;
                        if (line_length == 0) {
                            emptyLine = true; // flag empty lines, raise error only if this is embedded in the sequence
                        }
                    }
                } else {
                    entry.line_len = line_length + 1; // first line
                }
                entry.line_blen = entry.line_len - 1;
            }
            offset += line_length + 1;
        }
        // we've hit the end of the fasta file!
        // flush the last entry
        flushEntryToIndex(entry);
    } else {
        cerr << "could not open reference file " << refname << " for indexing!" << endl;
        exit(1);
    }
}

void FastaIndex::flushEntryToIndex(FastaIndexEntry& entry) {
    string name = split(entry.name, " \t").at(0);  // key by first token of name
    sequenceNames.push_back(name);
    this->insert(make_pair(name, FastaIndexEntry(entry.name, entry.length,
                        entry.offset, entry.line_blen,
                        entry.line_len)));

}

void FastaIndex::writeIndexFile(string fname) {
    //cerr << "writing fasta index file " << fname << endl;
    ofstream file;
    file.open(fname.c_str());
    if (file.is_open()) {
        file << *this;
    } else {
        cerr << "could not open index file " << fname << " for writing!" << endl;
        exit(1);
    }
}

FastaIndex::~FastaIndex(void) {
    indexFile.close();
}

FastaIndexEntry FastaIndex::entry(string name) {
    FastaIndex::iterator e = this->find(name);
    if (e == this->end()) {
        cerr << "unable to find FASTA index entry for '" << name << "'" << endl;
        exit(1);
    } else {
        return e->second;
    }
}

string FastaIndex::indexFileExtension() { return ".fai"; }

void FastaReference::open(string reffilename, bool usemmap) {
    filename = reffilename;
    if (!(file = fopen(filename.c_str(), "r"))) {
        cerr << "could not open " << filename << endl;
        exit(1);
    }
    index = new FastaIndex();
    struct stat stFileInfo;
    string indexFileName = filename + index->indexFileExtension();
    // if we can find an index file, use it
    if(stat(indexFileName.c_str(), &stFileInfo) == 0) {
        index->readIndexFile(indexFileName);
    } else { // otherwise, read the reference and generate the index file in the cwd
        cerr << "index file " << indexFileName << " not found, generating..." << endl;
        index->indexReference(filename);
        index->writeIndexFile(indexFileName);
    }
    if (usemmap) {
        usingmmap = true;
        int fd = fileno(file);
        struct stat sb;
        if (fstat(fd, &sb) == -1)
            cerr << "could not stat file" << filename << endl;
        filesize = sb.st_size;
        // map the whole file
        filemm = mmap(NULL, filesize, PROT_READ, MAP_SHARED, fd, 0);
    }
}

FastaReference::~FastaReference(void) {
    fclose(file);
    if (usingmmap) {
        munmap(filemm, filesize);
    }
    delete index;
}

string FastaReference::getSequence(string seqname) {
    FastaIndexEntry entry = index->entry(seqname);
    int newlines_in_sequence = entry.length / entry.line_blen;
    int seqlen = newlines_in_sequence  + entry.length;
    char* seq = (char*) calloc (seqlen + 1, sizeof(char));
    if (usingmmap) {
        memcpy(seq, (char*) filemm + entry.offset, seqlen);
    } else {
        fseek64(file, entry.offset, SEEK_SET);
        fread(seq, sizeof(char), seqlen, file);
    }
    seq[seqlen] = '\0';
    char* pbegin = seq;
    char* pend = seq + (seqlen/sizeof(char));
    pend = remove(pbegin, pend, '\n');
    pend = remove(pbegin, pend, '\0');
    string s = seq;
    free(seq);
    s.resize((pend - pbegin)/sizeof(char));
    return s;
}

// TODO cleanup; odd function.  use a map
string FastaReference::sequenceNameStartingWith(string seqnameStart) {
    try {
        return (*index)[seqnameStart].name;
    } catch (exception& e) {
        cerr << e.what() << ": unable to find index entry for " << seqnameStart << endl;
        exit(1);
    }
}

string FastaReference::getSubSequence(string seqname, int start, int length) {
    FastaIndexEntry entry = index->entry(seqname);
    if (start < 0 || length < 1) {
        cerr << "Error: cannot construct subsequence with negative offset or length < 1" << endl;
        exit(1);
    }
    // we have to handle newlines
    // approach: count newlines before start
    //           count newlines by end of read
    //             subtracting newlines before start find count of embedded newlines
    int newlines_before = start > 0 ? (start - 1) / entry.line_blen : 0;
    int newlines_by_end = (start + length - 1) / entry.line_blen;
    int newlines_inside = newlines_by_end - newlines_before;
    int seqlen = length + newlines_inside;
    char* seq = (char*) calloc (seqlen + 1, sizeof(char));
    if (usingmmap) {
        memcpy(seq, (char*) filemm + entry.offset + newlines_before + start, seqlen);
    } else {
        fseek64(file, (off_t) (entry.offset + newlines_before + start), SEEK_SET);
        fread(seq, sizeof(char), (off_t) seqlen, file);
    }
    seq[seqlen] = '\0';
    char* pbegin = seq;
    char* pend = seq + (seqlen/sizeof(char));
    pend = remove(pbegin, pend, '\n');
    pend = remove(pbegin, pend, '\0');
    string s = seq;
    free(seq);
    s.resize((pend - pbegin)/sizeof(char));
    return s;
}

long unsigned int FastaReference::sequenceLength(string seqname) {
    FastaIndexEntry entry = index->entry(seqname);
    return entry.length;
}

<commit_msg>tweak to Fasta.cpp<commit_after>// ***************************************************************************
// FastaIndex.cpp (c) 2010 Erik Garrison <erik.garrison@bc.edu>
// Marth Lab, Department of Biology, Boston College
// All rights reserved.
// ---------------------------------------------------------------------------
// Last modified: 9 February 2010 (EG)
// ---------------------------------------------------------------------------

#include "Fasta.h"

FastaIndexEntry::FastaIndexEntry(string name, int length, long long offset, int line_blen, int line_len)
    : name(name)
    , length(length)
    , offset(offset)
    , line_blen(line_blen)
    , line_len(line_len)
{}

FastaIndexEntry::FastaIndexEntry(void) // empty constructor
{ clear(); }

FastaIndexEntry::~FastaIndexEntry(void)
{}

void FastaIndexEntry::clear(void)
{
    name = "";
    length = NULL;
    offset = -1;  // no real offset will ever be below 0, so this allows us to
                  // check if we have already recorded a real offset
    line_blen = NULL;
    line_len  = NULL;
}

ostream& operator<<(ostream& output, const FastaIndexEntry& e) {
    // just write the first component of the name, for compliance with other tools
    output << split(e.name, ' ').at(0) << "\t" << e.length << "\t" << e.offset << "\t" <<
        e.line_blen << "\t" << e.line_len;
    return output;  // for multiple << operators.
}

FastaIndex::FastaIndex(void)
{}

void FastaIndex::readIndexFile(string fname) {
    string line;
    long long linenum = 0;
    indexFile.open(fname.c_str(), ifstream::in);
    if (indexFile.is_open()) {
        while (getline (indexFile, line)) {
            ++linenum;
            // the fai format defined in samtools is tab-delimited, every line being:
            // fai->name[i], (int)x.len, (long long)x.offset, (int)x.line_blen, (int)x.line_len
            vector<string> fields = split(line, '\t');
            if (fields.size() == 5) {  // if we don't get enough fields then there is a problem with the file
                // note that fields[0] is the sequence name
                char* end;
                string name = split(fields[0], " \t").at(0);  // key by first token of name
                sequenceNames.push_back(name);
                this->insert(make_pair(name, FastaIndexEntry(fields[0], atoi(fields[1].c_str()),
                                                    strtoll(fields[2].c_str(), &end, 10),
                                                    atoi(fields[3].c_str()),
                                                    atoi(fields[4].c_str()))));
            } else {
                cerr << "Warning: malformed fasta index file " << fname <<
                    "does not have enough fields @ line " << linenum << endl;
                cerr << line << endl;
                exit(1);
            }
        }
    } else {
        cerr << "could not open index file " << fname << endl;
        exit(1);
    }
}

// for consistency this should be a class method
bool fastaIndexEntryCompare ( FastaIndexEntry a, FastaIndexEntry b) { return (a.offset<b.offset); }

ostream& operator<<(ostream& output, FastaIndex& fastaIndex) {
    vector<FastaIndexEntry> sortedIndex;
    for(vector<string>::const_iterator it = fastaIndex.sequenceNames.begin(); it != fastaIndex.sequenceNames.end(); ++it)
    {
        sortedIndex.push_back(fastaIndex[*it]);
    }
    sort(sortedIndex.begin(), sortedIndex.end(), fastaIndexEntryCompare);
    for( vector<FastaIndexEntry>::iterator fit = sortedIndex.begin(); fit != sortedIndex.end(); ++fit) {
        output << *fit << endl;
    }
    return output;
}

void FastaIndex::indexReference(string refname) {
    // overview:
    //  for line in the reference fasta file
    //  track byte offset from the start of the file
    //  if line is a fasta header, take the name and dump the last sequnece to the index
    //  if line is a sequence, add it to the current sequence
    //cerr << "indexing fasta reference " << refname << endl;
    string line;
    FastaIndexEntry entry;  // an entry buffer used in processing
    entry.clear();
    int line_length = 0;
    long long offset = 0;  // byte offset from start of file
    long long line_number = 0; // current line number
    bool mismatchedLineLengths = false; // flag to indicate if our line length changes mid-file
                                        // this will be used to raise an error
                                        // if we have a line length change at
                                        // any line other than the last line in
                                        // the sequence
    bool emptyLine = false;  // flag to catch empty lines, which we allow for
                             // index generation only on the last line of the sequence
    ifstream refFile;
    refFile.open(refname.c_str());
    if (refFile.is_open()) {
        while (getline(refFile, line)) {
            ++line_number;
            line_length = line.length();
            if (line[0] == ';') {
                // fasta comment, skip
            } else if (line[0] == '+') {
                // fastq quality header
                getline(refFile, line);
                line_length = line.length();
                offset += line_length + 1;
                // get and don't handle the quality line
                getline(refFile, line);
                line_length = line.length();
            } else if (line[0] == '>' || line[0] == '@') { // fasta /fastq header
                // if we aren't on the first entry, push the last sequence into the index
                if (entry.name != "") {
                    mismatchedLineLengths = false; // reset line length error tracker for every new sequence
                    emptyLine = false;
                    flushEntryToIndex(entry);
                    entry.clear();
                }
                entry.name = line.substr(1, line_length - 1);
            } else { // we assume we have found a sequence line
                if (entry.offset == -1) // NB initially the offset is -1
                    entry.offset = offset;
                entry.length += line_length;
                if (entry.line_len) {
                    //entry.line_len = entry.line_len ? entry.line_len : line_length + 1;
                    if (mismatchedLineLengths || emptyLine) {
                        if (line_length == 0) {
                            emptyLine = true; // flag empty lines, raise error only if this is embedded in the sequence
                        } else {
                            if (emptyLine) {
                                cerr << "ERROR: embedded newline";
                            } else {
                                cerr << "ERROR: mismatched line lengths";
                            }
                            cerr << " at line " << line_number << " within sequence " << entry.name <<
                                endl << "File not suitable for fasta index generation." << endl;
                            exit(1);
                        }
                    }
                    // this flag is set here and checked on the next line
                    // because we may have reached the end of the sequence, in
                    // which case a mismatched line length is OK
                    if (entry.line_len != line_length + 1) {
                        mismatchedLineLengths = true;
                        if (line_length == 0) {
                            emptyLine = true; // flag empty lines, raise error only if this is embedded in the sequence
                        }
                    }
                } else {
                    entry.line_len = line_length + 1; // first line
                }
                entry.line_blen = entry.line_len - 1;
            }
            offset += line_length + 1;
        }
        // we've hit the end of the fasta file!
        // flush the last entry
        flushEntryToIndex(entry);
    } else {
        cerr << "could not open reference file " << refname << " for indexing!" << endl;
        exit(1);
    }
}

void FastaIndex::flushEntryToIndex(FastaIndexEntry& entry) {
    string name = split(entry.name, " \t").at(0);  // key by first token of name
    sequenceNames.push_back(name);
    this->insert(make_pair(name, FastaIndexEntry(entry.name, entry.length,
                        entry.offset, entry.line_blen,
                        entry.line_len)));

}

void FastaIndex::writeIndexFile(string fname) {
    //cerr << "writing fasta index file " << fname << endl;
    ofstream file;
    file.open(fname.c_str());
    if (file.is_open()) {
        file << *this;
    } else {
        cerr << "could not open index file " << fname << " for writing!" << endl;
        exit(1);
    }
}

FastaIndex::~FastaIndex(void) {
    indexFile.close();
}

FastaIndexEntry FastaIndex::entry(string name) {
    FastaIndex::iterator e = this->find(name);
    if (e == this->end()) {
        cerr << "unable to find FASTA index entry for '" << name << "'" << endl;
        exit(1);
    } else {
        return e->second;
    }
}

string FastaIndex::indexFileExtension() { return ".fai"; }

void FastaReference::open(string reffilename, bool usemmap) {
    filename = reffilename;
    if (!(file = fopen(filename.c_str(), "r"))) {
        cerr << "could not open " << filename << endl;
        exit(1);
    }
    index = new FastaIndex();
    struct stat stFileInfo;
    string indexFileName = filename + index->indexFileExtension();
    // if we can find an index file, use it
    if(stat(indexFileName.c_str(), &stFileInfo) == 0) {
        index->readIndexFile(indexFileName);
    } else { // otherwise, read the reference and generate the index file in the cwd
        cerr << "index file " << indexFileName << " not found, generating..." << endl;
        index->indexReference(filename);
        index->writeIndexFile(indexFileName);
    }
    if (usemmap) {
        usingmmap = true;
        int fd = fileno(file);
        struct stat sb;
        if (fstat(fd, &sb) == -1)
            cerr << "could not stat file" << filename << endl;
        filesize = sb.st_size;
        // map the whole file
        filemm = mmap(NULL, filesize, PROT_READ, MAP_SHARED, fd, 0);
    }
}

FastaReference::~FastaReference(void) {
    fclose(file);
    if (usingmmap) {
        munmap(filemm, filesize);
    }
    delete index;
}

string FastaReference::getSequence(string seqname) {
    FastaIndexEntry entry = index->entry(seqname);
    int newlines_in_sequence = entry.length / entry.line_blen;
    int seqlen = newlines_in_sequence  + entry.length;
    char* seq = (char*) calloc (seqlen + 1, sizeof(char));
    if (usingmmap) {
        memcpy(seq, (char*) filemm + entry.offset, seqlen);
    } else {
        fseek64(file, entry.offset, SEEK_SET);
        fread(seq, sizeof(char), seqlen, file);
    }
    seq[seqlen] = '\0';
    char* pbegin = seq;
    char* pend = seq + (seqlen/sizeof(char));
    pend = remove(pbegin, pend, '\n');
    pend = remove(pbegin, pend, '\0');
    string s = seq;
    free(seq);
    s.resize((pend - pbegin)/sizeof(char));
    return s;
}

// TODO cleanup; odd function.  use a map
string FastaReference::sequenceNameStartingWith(string seqnameStart) {
    try {
        return (*index)[seqnameStart].name;
    } catch (exception& e) {
        cerr << e.what() << ": unable to find index entry for " << seqnameStart << endl;
        exit(1);
    }
}

string FastaReference::getSubSequence(string seqname, int start, int length) {
    FastaIndexEntry entry = index->entry(seqname);
    if (start < 0 || length < 1) {
        cerr << "Error: cannot construct subsequence with negative offset or length < 1" << endl;
        exit(1);
    }
    // we have to handle newlines
    // approach: count newlines before start
    //           count newlines by end of read
    //             subtracting newlines before start find count of embedded newlines
    int newlines_before = start > 0 ? (start - 1) / entry.line_blen : 0;
    int newlines_by_end = (start + length - 1) / entry.line_blen;
    int newlines_inside = newlines_by_end - newlines_before;
    int seqlen = length + newlines_inside;
    char* seq = (char*) calloc (seqlen + 1, sizeof(char));
    if (usingmmap) {
        memcpy(seq, (char*) filemm + entry.offset + newlines_before + start, seqlen);
    } else {
        fseek64(file, (off_t) (entry.offset + newlines_before + start), SEEK_SET);
        fread(seq, sizeof(char), (off_t) seqlen, file);
    }
    seq[seqlen] = '\0';
    char* pbegin = seq;
    char* pend = seq + (seqlen/sizeof(char));
    pend = remove(pbegin, pend, '\n');
    pend = remove(pbegin, pend, '\0');
    string s = seq;
    free(seq);
    s.resize((pend - pbegin)/sizeof(char));
    return s;
}

long unsigned int FastaReference::sequenceLength(string seqname) {
    FastaIndexEntry entry = index->entry(seqname);
    return entry.length;
}

<|endoftext|>


<commit_before>#include "config.h"

#include "util.h"
#include <errno.h>
#include <string>
#include <cstdlib>
#include <time.h>
#include <librdkafka/rdkafkacpp.h>
#include "BroString.h"
#include "threading/SerialTypes.h"
#include "Kafka.h"
#include "kafkawriter.bif.h"

using namespace logging;
using namespace writer;
using threading::Value;
using threading::Field;
using namespace RdKafka;

class RandomPartitionerCallback : public RdKafka::PartitionerCb {
private:
	unsigned int seed;

public:
	int32_t partitioner_cb(const RdKafka::Topic *topic, const std::string *key,
							int32_t partition_count, void *msg_opaque)
	{
		return (int32_t)rand_r(&seed) % partition_count;
	}

	RandomPartitionerCallback(){
		seed = time(NULL);
	}
};

KafkaWriter::KafkaWriter(WriterFrontend* frontend) : WriterBackend(frontend)
{

	json_formatter = 0;
	//srand(time(NULL));
	producer = NULL;
	topic = NULL;

	// initialize kafka variables...
	broker_name_len = BifConst::KafkaLogger::broker_name->Len();
	broker_name = new char[broker_name_len + 1];
	memcpy(broker_name, BifConst::KafkaLogger::broker_name->Bytes(), broker_name_len);
	broker_name[broker_name_len] = 0;

    topic_name_len = BifConst::KafkaLogger::topic_name->Len();
    topic_name = new char[topic_name_len + 1];
    memcpy(topic_name, BifConst::KafkaLogger::topic_name->Bytes(), topic_name_len);
    topic_name[topic_name_len] = 0;

    compression_codec_len = BifConst::KafkaLogger::compression_codec->Len();
    compression_codec = new char[compression_codec_len + 1];
    memcpy(compression_codec, BifConst::KafkaLogger::compression_codec->Bytes(), compression_codec_len);
    compression_codec[compression_codec_len] = 0;

    // initialize varibles used to store extra data appended to every message
    // (sensor name and log type)
    int sensor_name_len = BifConst::KafkaLogger::sensor_name->Len();
    char* sensor_name = new char[sensor_name_len + 1];
    memcpy(sensor_name, BifConst::KafkaLogger::sensor_name->Bytes(), sensor_name_len);
    sensor_name[sensor_name_len] = 0;

    int type_name_len = strlen(Info().path);
    char* type_name = new char[type_name_len + 1];
    memcpy(type_name, Info().path, type_name_len);
    type_name[type_name_len] = 0;

    client_id_len = BifConst::KafkaLogger::client_id->Len() + strlen(Info().path) + 1;
    client_id = new char[client_id_len + 1];
    memcpy(client_id, BifConst::KafkaLogger::client_id->Bytes(), client_id_len);
    strcat(client_id, "-");
    strcat(client_id, type_name);
    client_id[client_id_len] = 0;

    json_formatter = new threading::formatter::AddingJSON(this,
    								threading::formatter::AddingJSON::TS_MILLIS,
    								sensor_name,
    								type_name,
    								BifConst::KafkaLogger::logstash_style_timestamp
    							   );
}

KafkaWriter::~KafkaWriter()
{
    delete [] broker_name;
    delete [] topic_name;
    delete [] client_id;
    delete [] compression_codec;
    delete [] fixed_fields;
    // I think I need to shut down the connection to the producer before deleting
    // these variables. Also, if I just blindly delete these two, bro segfaults when
    // shutting down. Confess I don't understand what's happening here.
    //delete producer;
    //delete topic;
    delete json_formatter;
}

threading::Field** KafkaWriter::MakeFields(const threading::Field* const* fields, int num_fields, std::string path){
	// create the renamed fields, based on user-supplied config.
	threading::Field** newFields = (threading::Field**)malloc(sizeof(threading::Field*) * (num_fields));

	// what I'd like to do is
	// first, grab the rename table for just this log
	// The config will have a table of table of strings.
	// the first table key is the name of the log (dns, http, etc)
	// the internal table key is the column name, the internal table value
	// will be the name to change that to.
	// loop over the existing fields, look up the field name in the rename table.
	// if it exists, create a new field entry with the new name, otherwise,
	// copy the existing field name in to the new field list.
	//
	// However, I can't get the bro TableVar Lookup to return anything
	// even for tables that I know have data in them. I'm clearly doing
	// something wrong. So,  hardcode the renames in the interest of
	// getting something done.
	for (int i = 0; i < num_fields; i++){
		std::string newName;

		if (strcmp(fields[i]->name, "ts") == 0)
		{
			newName = "timestamp";
		}
		else if (strcmp(fields[i]->name, "id.orig_h") == 0)
		{
			newName = "source_ip";
		}
		else if (strcmp(fields[i]->name, "id.orig_p") == 0)
		{
			newName = "source_port";
		}
		else if (strcmp(fields[i]->name, "id.resp_h") == 0)
		{
			newName = "dest_ip";
		}
		else if (strcmp(fields[i]->name,"id.resp_p") == 0)
		{
			newName = "dest_port";
		}

		if (newName.empty()){
			newFields[i] = new threading::Field(fields[i]->name,
												fields[i]->secondary_name,
												fields[i]->type,
												fields[i]->subtype,
												true);
		}
		else {
			newFields[i]= new threading::Field(newName.c_str(),
												fields[i]->secondary_name,
												fields[i]->type,
												fields[i]->subtype,
												true);
		}

	}

	return newFields;
}

bool KafkaWriter::DoInit(const WriterInfo& info, int num_fields, const threading::Field* const* fields)
{
	conf = RdKafka::Conf::create(RdKafka::Conf::CONF_GLOBAL);
	tconf = RdKafka::Conf::create(RdKafka::Conf::CONF_TOPIC);
	// note: hard-coding to use the random partitioner right now...
	RandomPartitionerCallback* part_cb = new RandomPartitionerCallback();
	std::string errstr;

	// set up kafka connection (get brokers, set partitioner, etc)
	if (conf->set("metadata.broker.list", broker_name, errstr) != RdKafka::Conf::CONF_OK){
		reporter->Error("Failed to set metatdata.broker.list: %s", errstr.c_str());
		return false;
	}
    int default_batch_size_len = BifConst::KafkaLogger::default_batch_size->Len();
    char* default_batch_size = new char[default_batch_size_len + 1];
    memcpy(default_batch_size, BifConst::KafkaLogger::default_batch_size->Bytes(), default_batch_size_len);
    default_batch_size[default_batch_size_len] = 0;

    int max_batch_size_len = BifConst::KafkaLogger::max_batch_size->Len();
    char* max_batch_size = new char[max_batch_size_len + 1];
    memcpy(max_batch_size, BifConst::KafkaLogger::max_batch_size->Bytes(), max_batch_size_len);
    max_batch_size[max_batch_size_len] = 0;

    int max_batch_interval_len = BifConst::KafkaLogger::max_batch_interval->Len();
    char* max_batch_interval = new char[max_batch_interval_len + 1];
    memcpy(max_batch_interval, BifConst::KafkaLogger::max_batch_interval->Bytes(), max_batch_interval_len);
    max_batch_interval[max_batch_interval_len] = 0;

	conf->set("compression.codec", compression_codec, errstr);
	conf->set("client.id", client_id, errstr);
	conf->set("batch.num.messages", default_batch_size, errstr);
	conf->set("queue.buffering.max.messages", max_batch_size, errstr);
	conf->set("queue.buffering.max.ms", max_batch_interval, errstr);
	conf->set("producer.type", "async", errstr);
	if (tconf->set("partitioner_cb",  part_cb, errstr) != RdKafka::Conf::CONF_OK){
		reporter->Error("failed to set partitioner for Kafka. %s", errstr.c_str());
	}

	producer = RdKafka::Producer::create(conf, errstr);
	if (!producer) {
		reporter->Error("Failed to create producer");
		return false;
	}
	topic = RdKafka::Topic::create(producer, topic_name, tconf, errstr);

	if (!topic) {
		reporter->Error("Failed to create topic.");
		return false;
	}

	// set up lookups and renamed fields.
	fixed_fields = MakeFields(fields, num_fields, Info().path);

	return true;
}

bool KafkaWriter::DoWrite(int num_fields, const Field* const * fields, Value** vals)
    {
	ODesc buffer;

    // this may look silly, but as of this writing, Kafka's default
    // partitioning is poor. if you do not supply a key, kafka will never
    // call your partition function, even if one is specified in the config.
    // What it will do instead is choose a partition at random when it starts
    // up, and send everything to that partition. So, you need to supply a
    // partition key if you want your partitioner to be used
    const std::string partition_key = "this is a key to trigger partitioning.";

	buffer.Clear();

    json_formatter->Describe(&buffer, num_fields, fixed_fields, vals);
    const char* bytes = (const char*)buffer.Bytes();
    std::string errstr;

	// actually send the data to Kafka.
    RdKafka::ErrorCode resp = producer->produce(topic,
    											RdKafka::Topic::PARTITION_UA,
    											RdKafka::Producer::MSG_COPY /* Copy payload */,
    											const_cast<char *>(bytes),
    											strlen(bytes),
    											&partition_key,
    											NULL);
    if (resp != RdKafka::ERR_NO_ERROR) {
        errstr = RdKafka::err2str(resp);
        reporter->Error("Produce failed: %s", errstr.c_str());
        reporter->Error("failed line: %s", bytes);
    }

    // Note: this bit here means that even if the send to kafka fails, we're just going to
    // drop the messages. Such is life.
    producer->poll(0);

    return true;
    }

bool KafkaWriter::DoSetBuf(bool enabled)
    {
    // Nothing to do.
    return true;
    }

bool KafkaWriter::DoFlush(double network_time)
    {
    // Nothing to do.
    return true;
    }

bool KafkaWriter::DoFinish(double network_time)
    {
    RdKafka::wait_destroyed(5000);
    return true;
    }

bool KafkaWriter::DoHeartbeat(double network_time, double current_time)
    {
	//nothing to do...all timing handled inside Kafka.
    return true;
    }

bool KafkaWriter::DoRotate(const char* rotated_path, double open, double close, bool terminating)
    {
    // Nothing to do.
    FinishedRotation();
    return true;
    }
<commit_msg>Adding new field renames.<commit_after>#include "config.h"

#include "util.h"
#include <errno.h>
#include <string>
#include <cstdlib>
#include <time.h>
#include <librdkafka/rdkafkacpp.h>
#include "BroString.h"
#include "threading/SerialTypes.h"
#include "Kafka.h"
#include "kafkawriter.bif.h"

using namespace logging;
using namespace writer;
using threading::Value;
using threading::Field;
using namespace RdKafka;

class RandomPartitionerCallback : public RdKafka::PartitionerCb {
private:
	unsigned int seed;

public:
	int32_t partitioner_cb(const RdKafka::Topic *topic, const std::string *key,
							int32_t partition_count, void *msg_opaque)
	{
		return (int32_t)rand_r(&seed) % partition_count;
	}

	RandomPartitionerCallback(){
		seed = time(NULL);
	}
};

KafkaWriter::KafkaWriter(WriterFrontend* frontend) : WriterBackend(frontend)
{

	json_formatter = 0;
	//srand(time(NULL));
	producer = NULL;
	topic = NULL;

	// initialize kafka variables...
	broker_name_len = BifConst::KafkaLogger::broker_name->Len();
	broker_name = new char[broker_name_len + 1];
	memcpy(broker_name, BifConst::KafkaLogger::broker_name->Bytes(), broker_name_len);
	broker_name[broker_name_len] = 0;

    topic_name_len = BifConst::KafkaLogger::topic_name->Len();
    topic_name = new char[topic_name_len + 1];
    memcpy(topic_name, BifConst::KafkaLogger::topic_name->Bytes(), topic_name_len);
    topic_name[topic_name_len] = 0;

    compression_codec_len = BifConst::KafkaLogger::compression_codec->Len();
    compression_codec = new char[compression_codec_len + 1];
    memcpy(compression_codec, BifConst::KafkaLogger::compression_codec->Bytes(), compression_codec_len);
    compression_codec[compression_codec_len] = 0;

    // initialize varibles used to store extra data appended to every message
    // (sensor name and log type)
    int sensor_name_len = BifConst::KafkaLogger::sensor_name->Len();
    char* sensor_name = new char[sensor_name_len + 1];
    memcpy(sensor_name, BifConst::KafkaLogger::sensor_name->Bytes(), sensor_name_len);
    sensor_name[sensor_name_len] = 0;

    int type_name_len = strlen(Info().path);
    char* type_name = new char[type_name_len + 1];
    memcpy(type_name, Info().path, type_name_len);
    type_name[type_name_len] = 0;

    client_id_len = BifConst::KafkaLogger::client_id->Len() + strlen(Info().path) + 1;
    client_id = new char[client_id_len + 1];
    memcpy(client_id, BifConst::KafkaLogger::client_id->Bytes(), client_id_len);
    strcat(client_id, "-");
    strcat(client_id, type_name);
    client_id[client_id_len] = 0;

    json_formatter = new threading::formatter::AddingJSON(this,
    								threading::formatter::AddingJSON::TS_MILLIS,
    								sensor_name,
    								type_name,
    								BifConst::KafkaLogger::logstash_style_timestamp
    							   );
}

KafkaWriter::~KafkaWriter()
{
    delete [] broker_name;
    delete [] topic_name;
    delete [] client_id;
    delete [] compression_codec;
    delete [] fixed_fields;
    // I think I need to shut down the connection to the producer before deleting
    // these variables. Also, if I just blindly delete these two, bro segfaults when
    // shutting down. Confess I don't understand what's happening here.
    //delete producer;
    //delete topic;
    delete json_formatter;
}

threading::Field** KafkaWriter::MakeFields(const threading::Field* const* fields, int num_fields, std::string path){
	// create the renamed fields, based on user-supplied config.
	threading::Field** newFields = (threading::Field**)malloc(sizeof(threading::Field*) * (num_fields));

	// what I'd like to do is
	// first, grab the rename table for just this log
	// The config will have a table of table of strings.
	// the first table key is the name of the log (dns, http, etc)
	// the internal table key is the column name, the internal table value
	// will be the name to change that to.
	// loop over the existing fields, look up the field name in the rename table.
	// if it exists, create a new field entry with the new name, otherwise,
	// copy the existing field name in to the new field list.
	//
	// However, I can't get the bro TableVar Lookup to return anything
	// even for tables that I know have data in them. I'm clearly doing
	// something wrong. So,  hardcode the renames in the interest of
	// getting something done.
	//
	// Also, need to remove "."s from names for ElasticSearch.
	//
	for (int i = 0; i < num_fields; i++){
		std::string newName;

		if (strcmp(fields[i]->name, "ts") == 0)
		{
			newName = "timestamp";
		}
		else if (strcmp(fields[i]->name, "id.orig_h") == 0)
		{
			newName = "source_ip";
		}
		else if (strcmp(fields[i]->name, "id.orig_p") == 0)
		{
			newName = "source_port";
		}
		else if (strcmp(fields[i]->name, "id.resp_h") == 0)
		{
			newName = "dest_ip";
		}
		else if (strcmp(fields[i]->name,"id.resp_p") == 0)
		{
			newName = "dest_port";
		}
		else if (strcmp(fields[i]->name, "seen.indicator") == 0)
		{
			newName = "indicator";
		}
		else if (strcmp(fields[i]->name, "seen.indicator_type") == 0)
		{
			newName = "indicator_type";
		}
		else if (strcmp(fields[i]->name, "seen.where") == 0)
		{
			newName = "seen_where";
		}
		else if (strcmp(fields[i]->name, "seen.node") == 0)
		{
			newName = "seen_node";
		}

		if (newName.empty()){
			newFields[i] = new threading::Field(fields[i]->name,
												fields[i]->secondary_name,
												fields[i]->type,
												fields[i]->subtype,
												true);
		}
		else {
			newFields[i]= new threading::Field(newName.c_str(),
												fields[i]->secondary_name,
												fields[i]->type,
												fields[i]->subtype,
												true);
		}

	}

	return newFields;
}

bool KafkaWriter::DoInit(const WriterInfo& info, int num_fields, const threading::Field* const* fields)
{
	conf = RdKafka::Conf::create(RdKafka::Conf::CONF_GLOBAL);
	tconf = RdKafka::Conf::create(RdKafka::Conf::CONF_TOPIC);
	// note: hard-coding to use the random partitioner right now...
	RandomPartitionerCallback* part_cb = new RandomPartitionerCallback();
	std::string errstr;

	// set up kafka connection (get brokers, set partitioner, etc)
	if (conf->set("metadata.broker.list", broker_name, errstr) != RdKafka::Conf::CONF_OK){
		reporter->Error("Failed to set metatdata.broker.list: %s", errstr.c_str());
		return false;
	}
    int default_batch_size_len = BifConst::KafkaLogger::default_batch_size->Len();
    char* default_batch_size = new char[default_batch_size_len + 1];
    memcpy(default_batch_size, BifConst::KafkaLogger::default_batch_size->Bytes(), default_batch_size_len);
    default_batch_size[default_batch_size_len] = 0;

    int max_batch_size_len = BifConst::KafkaLogger::max_batch_size->Len();
    char* max_batch_size = new char[max_batch_size_len + 1];
    memcpy(max_batch_size, BifConst::KafkaLogger::max_batch_size->Bytes(), max_batch_size_len);
    max_batch_size[max_batch_size_len] = 0;

    int max_batch_interval_len = BifConst::KafkaLogger::max_batch_interval->Len();
    char* max_batch_interval = new char[max_batch_interval_len + 1];
    memcpy(max_batch_interval, BifConst::KafkaLogger::max_batch_interval->Bytes(), max_batch_interval_len);
    max_batch_interval[max_batch_interval_len] = 0;

	conf->set("compression.codec", compression_codec, errstr);
	conf->set("client.id", client_id, errstr);
	conf->set("batch.num.messages", default_batch_size, errstr);
	conf->set("queue.buffering.max.messages", max_batch_size, errstr);
	conf->set("queue.buffering.max.ms", max_batch_interval, errstr);
	conf->set("producer.type", "async", errstr);
	if (tconf->set("partitioner_cb",  part_cb, errstr) != RdKafka::Conf::CONF_OK){
		reporter->Error("failed to set partitioner for Kafka. %s", errstr.c_str());
	}

	producer = RdKafka::Producer::create(conf, errstr);
	if (!producer) {
		reporter->Error("Failed to create producer");
		return false;
	}
	topic = RdKafka::Topic::create(producer, topic_name, tconf, errstr);

	if (!topic) {
		reporter->Error("Failed to create topic.");
		return false;
	}

	// set up lookups and renamed fields.
	fixed_fields = MakeFields(fields, num_fields, Info().path);

	return true;
}

bool KafkaWriter::DoWrite(int num_fields, const Field* const * fields, Value** vals)
    {
	ODesc buffer;

    // this may look silly, but as of this writing, Kafka's default
    // partitioning is poor. if you do not supply a key, kafka will never
    // call your partition function, even if one is specified in the config.
    // What it will do instead is choose a partition at random when it starts
    // up, and send everything to that partition. So, you need to supply a
    // partition key if you want your partitioner to be used
    const std::string partition_key = "this is a key to trigger partitioning.";

	buffer.Clear();

    json_formatter->Describe(&buffer, num_fields, fixed_fields, vals);
    const char* bytes = (const char*)buffer.Bytes();
    std::string errstr;

	// actually send the data to Kafka.
    RdKafka::ErrorCode resp = producer->produce(topic,
    											RdKafka::Topic::PARTITION_UA,
    											RdKafka::Producer::MSG_COPY /* Copy payload */,
    											const_cast<char *>(bytes),
    											strlen(bytes),
    											&partition_key,
    											NULL);
    if (resp != RdKafka::ERR_NO_ERROR) {
        errstr = RdKafka::err2str(resp);
        reporter->Error("Produce failed: %s", errstr.c_str());
        reporter->Error("failed line: %s", bytes);
    }

    // Note: this bit here means that even if the send to kafka fails, we're just going to
    // drop the messages. Such is life.
    producer->poll(0);

    return true;
    }

bool KafkaWriter::DoSetBuf(bool enabled)
    {
    // Nothing to do.
    return true;
    }

bool KafkaWriter::DoFlush(double network_time)
    {
    // Nothing to do.
    return true;
    }

bool KafkaWriter::DoFinish(double network_time)
    {
    RdKafka::wait_destroyed(5000);
    return true;
    }

bool KafkaWriter::DoHeartbeat(double network_time, double current_time)
    {
	//nothing to do...all timing handled inside Kafka.
    return true;
    }

bool KafkaWriter::DoRotate(const char* rotated_path, double open, double close, bool terminating)
    {
    // Nothing to do.
    FinishedRotation();
    return true;
    }
<|endoftext|>


<commit_before>#include "log.h"
#include <iostream>

namespace share_me_utils {

Log::Log() {
  memset(m_prefixSwitchs, -1, sizeof(m_prefixSwitchs));
  m_prefixSymbols[eDate] = 'd';
  m_prefixSymbols[eTime] = 't';
  m_prefixSymbols[eFile] = 'F';
  m_prefixSymbols[eFunc] = 'f';
  m_prefixSymbols[eLine] = 'l';

  m_levelString[S_TRACE] = "TRACE";
  m_levelString[S_DEBUG] = "DEBUG";
  m_levelString[S_INFO] = "INFO";
  m_levelString[S_WARN] = "WARN";
  m_levelString[S_ERROR] = "ERROR";
  m_levelString[S_FATAL] = "FATAL";
  SetPrefix("%d %t %F %f %l");
  initColor();
#ifdef _WIN32
  m_logEvent = CreateEvent(NULL, FALSE, FALSE, NULL);
#endif
}

bool Log::initColor() {
  memset(m_levelColor, 0, sizeof(m_levelColor));
  HANDLE stdHandle = GetStdHandle(STD_OUTPUT_HANDLE);
  CONSOLE_SCREEN_BUFFER_INFO csbiInfo;
  GetConsoleScreenBufferInfo(stdHandle, &csbiInfo);
  m_oldColorAttr = csbiInfo.wAttributes;
  m_levelColor[S_TRACE] = FOREGROUND_BLUE | FOREGROUND_INTENSITY;
  m_levelColor[S_DEBUG] = m_oldColorAttr | FOREGROUND_INTENSITY;
  m_levelColor[S_INFO] = FOREGROUND_GREEN | FOREGROUND_INTENSITY;
  m_levelColor[S_WARN] =
      FOREGROUND_GREEN | FOREGROUND_RED | FOREGROUND_INTENSITY;
  m_levelColor[S_ERROR] = FOREGROUND_RED | FOREGROUND_INTENSITY;
  m_levelColor[S_FATAL] =
      FOREGROUND_RED | FOREGROUND_BLUE | FOREGROUND_INTENSITY;
  return true;
}

void Log::setColor(int level) {
  static HANDLE stdHandle = NULL;
  if (level >= S_TRACE && level < S_INVALID) {
    if (!stdHandle) {
      stdHandle = GetStdHandle(STD_OUTPUT_HANDLE);
    }
    SetConsoleTextAttribute(stdHandle, m_levelColor[level]);
  }
}

void Log::resetColor() {
  static HANDLE stdHandle = NULL;
  if (!stdHandle) {
    stdHandle = GetStdHandle(STD_OUTPUT_HANDLE);
  }
  SetConsoleTextAttribute(stdHandle, m_oldColorAttr);
}

Log *Log::Instance() {
  static Log *inst = NULL;
  if (!inst)
    inst = new Log();
  return inst;
}

bool Log::Start() {
  DWORD threadID;
  HANDLE threadHandle;
  threadHandle =
      CreateThread(NULL, 0, Log::Instance()->loop, nullptr, 0, &threadID);
  return threadHandle != NULL;
}

bool Log::SetPrefix(const char *prefix) {
  if (!prefix)
    return false;
  if (*prefix == '\0')
    return false;

  const char *s = prefix;
  bool switchConfiged[eTop] = {false};
  memset(m_prefixSwitchs, -1, sizeof(m_prefixSwitchs));
  for (int i = 0; s[i] != '\0'; ++i) {
    if (s[i] == '%') {
      if (i >= 1 && s[i - 1] == '\\')
        continue;
      for (int j = eDate; j < eTop; ++j) {
        if (m_prefixSymbols[j] == s[i + 1]) {
          if (!switchConfiged[j]) {
            for (int k = eDate; k < eTop; ++k) {
              if (m_prefixSwitchs[k] == -1) {
                m_prefixSwitchs[k] = j;
                break;
              }
            }
          }
          break;
        }
      }
    }
  }
  return true;
}

void Log::LogContent(const char *filename, const int lineno,
                     const char *funcname, int level, const char *format, ...) {
  int offset = 0;
  int prefixLen = 0;
  prefixLen = generatePrefix(filename, funcname, lineno, level);
  if (prefixLen > 0) {
    // m_prefixBuffer[offset] = '\0';
    memcpy(m_logBuffer, m_prefixBuffer, prefixLen);
  } else
    prefixLen = 0;
  va_list args; //定义一个va_list类型的变量，用来储存单个参数
  va_start(args, format); //使args指向可变参数的第一个参数
  offset = vsnprintf(m_logBuffer + prefixLen, 2 * LOG_BUFFER_LENGTH - prefixLen,
                     format, args); //必须用vprintf等带V的
  va_end(args);
  if (offset > 0) {
    offset = offset + prefixLen >= 2 * LOG_BUFFER_LENGTH - 2
                 ? (2 * LOG_BUFFER_LENGTH - 2)
                 : offset + prefixLen;
  } else {
    offset = 0;
  }
  m_logBuffer[offset] = '\n';
  m_logBuffer[offset + 1] = '\0';
  setColor(level);
  fprintf(stdout, "%s", m_logBuffer);
  resetColor();
}

int Log::generatePrefix(const char *filename, const char *funcname,
                        const int lineno, int level) {
  memset(m_prefixBuffer, 0, LOG_BUFFER_LENGTH);
  size_t offset = 0;
  for (int i = eDate; i < eTop; ++i) {
    if (m_prefixSwitchs[i] == -1)
      break;
    switch (m_prefixSwitchs[i]) {
    case eDate: {
      time_t rawtime;
      struct tm timeinfo;
      time(&rawtime);
      localtime(&timeinfo, &rawtime);
      offset += strftime(m_prefixBuffer + offset, LOG_BUFFER_LENGTH - offset,
                         "[%Y-%m-%d]", &timeinfo);
      // offset += sprintf(m_prefixBuffer + offset, "%s", date);
      break;
    }
    case eTime: {
      time_t rawtime;
      struct tm timeinfo;
      time(&rawtime);
      localtime(&timeinfo, &rawtime);
      offset += strftime(m_prefixBuffer + offset, LOG_BUFFER_LENGTH - offset,
                         "[%H:%M:%S]", &timeinfo);
      // offset += sprintf(m_prefixBuffer + offset, "%s", time);
      break;
    }
    case eFile: {
      offset += snprintf(m_prefixBuffer + offset,
                         2 * LOG_BUFFER_LENGTH - offset, "[%s]", filename);
      break;
    }
    case eFunc: {
      offset += snprintf(m_prefixBuffer + offset,
                         2 * LOG_BUFFER_LENGTH - offset, "[%s]", funcname);
      break;
    }
    case eLine: {
      offset += snprintf(m_prefixBuffer + offset,
                         2 * LOG_BUFFER_LENGTH - offset, "[%d]", lineno);
      break;
    }
    default: { break; }
    }
  }
  offset += snprintf(m_prefixBuffer + offset, 2 * LOG_BUFFER_LENGTH - offset,
                     "[%s]", m_levelString[level]);
  return (offset > 0 && offset < LOG_BUFFER_LENGTH) ? (int)offset : -1;
}

void Log::loop(THREAD_PARAM parma) {
  MsgNode* msgs = nullptr;
  MsgNode* msg = nullptr;
  while (true) {
    msgs = m_msgQueue.get();
    if (!msgs) {
      DWORD dReturn = WaitForSingleObject(m_logEvent, 100)
      switch (dReturn) {
        case WAIT_TIMEOUT:
        case WAIT_OBJECT_0: continue;
        case WAIT_ABANDONED: break;
        case WAIT_FAILED: break;
      }
    }
    while(msgs) {
      msg = msgs;
      out(msg->msg);
      msgs = msgs->next;
      DESTROY_MSG_NODE(msg);
    }
  }
}

bool Log::AppendMsg(LogMsg *msg) { return MsgQueue.Append(msg); }

void Log::out(LogMsg* msg) {
  int offset = 0;
  offset = generatePrefix(msg->fileName, msg->funcName, msg->lineno, msg->logLevel);
  if (offset > 0) {
    memcpy(m_logBuffer, m_prefixBuffer, offset);
  } else
    offset = 0;
  offset += sprintf(m_logBuffer+offset, 2 * LOG_BUFFER_LENGTH-offset, "%s", msg->msg);
  if (offset > 0) {
    offset = offset >= 2 * LOG_BUFFER_LENGTH - 2
                 ? (2 * LOG_BUFFER_LENGTH - 2)
                 : offset;
  } else {
    offset = 0;
  }
  m_logBuffer[offset] = '\n';
  m_logBuffer[offset + 1] = '\0';
  setColor(msg->logLevel);
  fprintf(stdout, "%s", m_logBuffer);
  resetColor();
}

void Log::Notify() {
  if (m_isRunning) return;
  int res = SetEvent(m_logEvent);
}

Log::MsgQueue::MsgQueue() : m_head(nullptr), m_tail(nullptr), m_count(0) {}
Log::MsgQueue::~MsgQueue() {
  MsgNode* node = nullptr;
  while(m_head) {
    node = m_head;
    m_head = m_head->next;
    delete node->msg;
    delete node;
  }
  m_count = 0;
}
bool Log::MsgQueue::Append(LogMsg *msg) {
  if (m_count > MAX_COUNT) {
    return false;
  }
  MsgNode *node = new MsgNode;
  node->msg = msg;
  node->next = nullptr;
  if (!m_head) {
    m_head = node;
  } else {
    m_tail->next = node;
  }
  m_tail = node;
  ++m_count;
  return true;
}
MsgNode *Log::MsgQueue::get() {
  MsgNode *node = m_head;
  m_head = nullptr;
  m_tail = nullptr;
  m_count = 0;
  return node;
}

void Log::MsgQueue::Stop() {
  m_count = MAX_COUNT + 1;
}

void Log::formatString(const char *format, ...) { ; }

// Logger implement

Logger::Logger() {}
Logger::~Logger() {}
void Logger::SendLog(const char *filename, const int lineno,
                     const char *funcname, int level, const char *format, ...) {
  int offset = 0;
  char *logBuffer = new char[LOG_BUFFER_LENGTH];
  va_list args; //定义一个va_list类型的变量，用来储存单个参数
  va_start(args, format); //使args指向可变参数的第一个参数
  offset = vsnprintf(logBuffer, LOG_BUFFER_LENGTH, format,
                     args); //必须用vprintf等带V的
  va_end(args);
  if (offset > 0) {
    offset = offset + prefixLen >= LOG_BUFFER_LENGTH - 2
                 ? (LOG_BUFFER_LENGTH - 2)
                 : offset;
  } else {
    offset = 0;
  }
  logBuffer[offset] = '\n';
  logBuffer[offset + 1] = '\0';
  LogMsg msg = new LogMsg(filename, lineno, funcname, level, logBuffer);
  Log::Instance()->AppendMsg(msg);
}
} // namespace share_me_utils<commit_msg>event<commit_after>#include "log.h"
#include <iostream>

namespace share_me_utils {

Log::Log() {
  memset(m_prefixSwitchs, -1, sizeof(m_prefixSwitchs));
  m_prefixSymbols[eDate] = 'd';
  m_prefixSymbols[eTime] = 't';
  m_prefixSymbols[eFile] = 'F';
  m_prefixSymbols[eFunc] = 'f';
  m_prefixSymbols[eLine] = 'l';

  m_levelString[S_TRACE] = "TRACE";
  m_levelString[S_DEBUG] = "DEBUG";
  m_levelString[S_INFO] = "INFO";
  m_levelString[S_WARN] = "WARN";
  m_levelString[S_ERROR] = "ERROR";
  m_levelString[S_FATAL] = "FATAL";
  SetPrefix("%d %t %F %f %l");
  initColor();
#ifdef _WIN32
  m_logEvent = CreateEvent(NULL, FALSE, FALSE, NULL);
#endif
}

bool Log::initColor() {
  memset(m_levelColor, 0, sizeof(m_levelColor));
  HANDLE stdHandle = GetStdHandle(STD_OUTPUT_HANDLE);
  CONSOLE_SCREEN_BUFFER_INFO csbiInfo;
  GetConsoleScreenBufferInfo(stdHandle, &csbiInfo);
  m_oldColorAttr = csbiInfo.wAttributes;
  m_levelColor[S_TRACE] = FOREGROUND_BLUE | FOREGROUND_INTENSITY;
  m_levelColor[S_DEBUG] = m_oldColorAttr | FOREGROUND_INTENSITY;
  m_levelColor[S_INFO] = FOREGROUND_GREEN | FOREGROUND_INTENSITY;
  m_levelColor[S_WARN] =
      FOREGROUND_GREEN | FOREGROUND_RED | FOREGROUND_INTENSITY;
  m_levelColor[S_ERROR] = FOREGROUND_RED | FOREGROUND_INTENSITY;
  m_levelColor[S_FATAL] =
      FOREGROUND_RED | FOREGROUND_BLUE | FOREGROUND_INTENSITY;
  return true;
}

void Log::setColor(int level) {
  static HANDLE stdHandle = NULL;
  if (level >= S_TRACE && level < S_INVALID) {
    if (!stdHandle) {
      stdHandle = GetStdHandle(STD_OUTPUT_HANDLE);
    }
    SetConsoleTextAttribute(stdHandle, m_levelColor[level]);
  }
}

void Log::resetColor() {
  static HANDLE stdHandle = NULL;
  if (!stdHandle) {
    stdHandle = GetStdHandle(STD_OUTPUT_HANDLE);
  }
  SetConsoleTextAttribute(stdHandle, m_oldColorAttr);
}

Log *Log::Instance() {
  static Log *inst = NULL;
  if (!inst)
    inst = new Log();
  return inst;
}

bool Log::Start() {
  DWORD threadID;
  HANDLE threadHandle;
  threadHandle =
      CreateThread(NULL, 0, Log::Instance()->loop, nullptr, 0, &threadID);
  return threadHandle != NULL;
}

bool Log::SetPrefix(const char *prefix) {
  if (!prefix)
    return false;
  if (*prefix == '\0')
    return false;

  const char *s = prefix;
  bool switchConfiged[eTop] = {false};
  memset(m_prefixSwitchs, -1, sizeof(m_prefixSwitchs));
  for (int i = 0; s[i] != '\0'; ++i) {
    if (s[i] == '%') {
      if (i >= 1 && s[i - 1] == '\\')
        continue;
      for (int j = eDate; j < eTop; ++j) {
        if (m_prefixSymbols[j] == s[i + 1]) {
          if (!switchConfiged[j]) {
            for (int k = eDate; k < eTop; ++k) {
              if (m_prefixSwitchs[k] == -1) {
                m_prefixSwitchs[k] = j;
                break;
              }
            }
          }
          break;
        }
      }
    }
  }
  return true;
}

void Log::LogContent(const char *filename, const int lineno,
                     const char *funcname, int level, const char *format, ...) {
  int offset = 0;
  int prefixLen = 0;
  prefixLen = generatePrefix(filename, funcname, lineno, level);
  if (prefixLen > 0) {
    // m_prefixBuffer[offset] = '\0';
    memcpy(m_logBuffer, m_prefixBuffer, prefixLen);
  } else
    prefixLen = 0;
  va_list args; //定义一个va_list类型的变量，用来储存单个参数
  va_start(args, format); //使args指向可变参数的第一个参数
  offset = vsnprintf(m_logBuffer + prefixLen, 2 * LOG_BUFFER_LENGTH - prefixLen,
                     format, args); //必须用vprintf等带V的
  va_end(args);
  if (offset > 0) {
    offset = offset + prefixLen >= 2 * LOG_BUFFER_LENGTH - 2
                 ? (2 * LOG_BUFFER_LENGTH - 2)
                 : offset + prefixLen;
  } else {
    offset = 0;
  }
  m_logBuffer[offset] = '\n';
  m_logBuffer[offset + 1] = '\0';
  setColor(level);
  fprintf(stdout, "%s", m_logBuffer);
  resetColor();
}

int Log::generatePrefix(const char *filename, const char *funcname,
                        const int lineno, int level) {
  memset(m_prefixBuffer, 0, LOG_BUFFER_LENGTH);
  size_t offset = 0;
  for (int i = eDate; i < eTop; ++i) {
    if (m_prefixSwitchs[i] == -1)
      break;
    switch (m_prefixSwitchs[i]) {
    case eDate: {
      time_t rawtime;
      struct tm timeinfo;
      time(&rawtime);
      localtime(&timeinfo, &rawtime);
      offset += strftime(m_prefixBuffer + offset, LOG_BUFFER_LENGTH - offset,
                         "[%Y-%m-%d]", &timeinfo);
      // offset += sprintf(m_prefixBuffer + offset, "%s", date);
      break;
    }
    case eTime: {
      time_t rawtime;
      struct tm timeinfo;
      time(&rawtime);
      localtime(&timeinfo, &rawtime);
      offset += strftime(m_prefixBuffer + offset, LOG_BUFFER_LENGTH - offset,
                         "[%H:%M:%S]", &timeinfo);
      // offset += sprintf(m_prefixBuffer + offset, "%s", time);
      break;
    }
    case eFile: {
      offset += snprintf(m_prefixBuffer + offset,
                         2 * LOG_BUFFER_LENGTH - offset, "[%s]", filename);
      break;
    }
    case eFunc: {
      offset += snprintf(m_prefixBuffer + offset,
                         2 * LOG_BUFFER_LENGTH - offset, "[%s]", funcname);
      break;
    }
    case eLine: {
      offset += snprintf(m_prefixBuffer + offset,
                         2 * LOG_BUFFER_LENGTH - offset, "[%d]", lineno);
      break;
    }
    default: { break; }
    }
  }
  offset += snprintf(m_prefixBuffer + offset, 2 * LOG_BUFFER_LENGTH - offset,
                     "[%s]", m_levelString[level]);
  return (offset > 0 && offset < LOG_BUFFER_LENGTH) ? (int)offset : -1;
}

void Log::loop(THREAD_PARAM parma) {
  MsgNode* msgs = nullptr;
  MsgNode* msg = nullptr;
  while (true) {
    msgs = m_msgQueue.get();
    if (!msgs) {
      DWORD dReturn = WaitForSingleObject(m_logEvent, 100)
      switch (dReturn) {
        case WAIT_TIMEOUT:
        case WAIT_OBJECT_0: continue;
        case WAIT_ABANDONED: break;
        case WAIT_FAILED: break;
      }
    }
    while(msgs) {
      msg = msgs;
      out(msg->msg);
      msgs = msgs->next;
      DESTROY_MSG_NODE(msg);
    }
  }
}

bool Log::AppendMsg(LogMsg *msg) { return MsgQueue.Append(msg); }

void Log::out(LogMsg* msg) {
  int offset = 0;
  offset = generatePrefix(msg->fileName, msg->funcName, msg->lineno, msg->logLevel);
  if (offset > 0) {
    memcpy(m_logBuffer, m_prefixBuffer, offset);
  } else
    offset = 0;
  offset += sprintf(m_logBuffer+offset, 2 * LOG_BUFFER_LENGTH-offset, "%s", msg->msg);
  if (offset > 0) {
    offset = offset >= 2 * LOG_BUFFER_LENGTH - 2
                 ? (2 * LOG_BUFFER_LENGTH - 2)
                 : offset;
  } else {
    offset = 0;
  }
  m_logBuffer[offset] = '\n';
  m_logBuffer[offset + 1] = '\0';
  setColor(msg->logLevel);
  fprintf(stdout, "%s", m_logBuffer);
  resetColor();
}

void Log::Notify() {
  if (m_isRunning) return;
  SetEvent(m_logEvent);
}

Log::MsgQueue::MsgQueue() : m_head(nullptr), m_tail(nullptr), m_count(0) {}
Log::MsgQueue::~MsgQueue() {
  MsgNode* node = nullptr;
  while(m_head) {
    node = m_head;
    m_head = m_head->next;
    delete node->msg;
    delete node;
  }
  m_count = 0;
}
bool Log::MsgQueue::Append(LogMsg *msg) {
  if (m_count > MAX_COUNT) {
    return false;
  }
  MsgNode *node = new MsgNode;
  node->msg = msg;
  node->next = nullptr;
  if (!m_head) {
    m_head = node;
  } else {
    m_tail->next = node;
  }
  m_tail = node;
  ++m_count;
  return true;
}
MsgNode *Log::MsgQueue::get() {
  MsgNode *node = m_head;
  m_head = nullptr;
  m_tail = nullptr;
  m_count = 0;
  return node;
}

void Log::MsgQueue::Stop() {
  m_count = MAX_COUNT + 1;
}

void Log::formatString(const char *format, ...) { ; }

// Logger implement

Logger::Logger() {}
Logger::~Logger() {}
void Logger::SendLog(const char *filename, const int lineno,
                     const char *funcname, int level, const char *format, ...) {
  int offset = 0;
  char *logBuffer = new char[LOG_BUFFER_LENGTH];
  va_list args; //定义一个va_list类型的变量，用来储存单个参数
  va_start(args, format); //使args指向可变参数的第一个参数
  offset = vsnprintf(logBuffer, LOG_BUFFER_LENGTH, format,
                     args); //必须用vprintf等带V的
  va_end(args);
  if (offset > 0) {
    offset = offset + prefixLen >= LOG_BUFFER_LENGTH - 2
                 ? (LOG_BUFFER_LENGTH - 2)
                 : offset;
  } else {
    offset = 0;
  }
  logBuffer[offset] = '\n';
  logBuffer[offset + 1] = '\0';
  LogMsg msg = new LogMsg(filename, lineno, funcname, level, logBuffer);
  Log::Instance()->AppendMsg(msg);
}
} // namespace share_me_utils<|endoftext|>


<commit_before>#include <iostream>
#include <fstream>
#include <stdexcept>
#include <string>
#include <vector>
#include <algorithm>
#include <functional>
#include <unordered_map>

#include "global.hpp"
#include "operators.hpp"
#include "tokens.hpp"
#include "nodes.hpp"
#include "builtins.hpp"

namespace lang {

  class Parser {
  private:
    std::vector<Token> tokens {};
    inline void skipCharacters(unsigned int& i, int by) {i += by;}
    // If we're already at the next character, but the loop will increment i by 1
    inline void preventIncrement(unsigned int& i) {i--;}

    std::vector<Token> variables {};
    std::vector<std::string> keywords {"define", "if", "while"};
    std::vector<std::string> constructKeywords {"do", "end", "else"};
  public:
    AST tree = AST();

    Parser(std::string code) {
      if (PARSER_PRINT_INPUT) print(code, "\n");
      try {
        tokenize(code);
        if (PARSER_PRINT_TOKENS) for (auto tok : tokens) print(tok, "\n");
        if (PARSER_PRINT_AS_EXPR) for (auto tok : ExpressionNode(tokens).getRPNOutput()) print(tok.data, " ");
        buildTree(tokens);
      } catch (SyntaxError &e) {
        print(e.toString(), "\n");
      }
    }

    void tokenize(std::string code) {
      unsigned int lines = 0;
      for (unsigned int i = 0; i < code.length(); ++i) {
        if (code[i] == '\n') lines++;

        // Ignore whitespace
        if (isspace(code[i])) continue;

        // Block begin/end
        // Parenthesis
        // Semicolons
        if (code[i] == '{' || code[i] == '}' ||
        code[i] == '(' || code[i] == ')' ||
        code[i] == ';') {
          tokens.push_back(Token(std::string(1, code[i]), CONSTRUCT, lines));
          continue;
        }

        // Operators
        auto initTokSize = tokens.size();
        std::for_each(opList.begin(), opList.end(), [this, initTokSize, code, i, lines](Operator& op) {
          if (initTokSize < tokens.size()) return; // If there are more tokens than before for_each, the operator was already added, so return.
          if (op.toString().length() + i > code.length()) return; // If the operator is longer than the source string, ignore it.
          if (op.toString() == code.substr(i, op.toString().length())) {
            Operator* tmp = &op;
            // TODO: apply DRY on these ifs
            if (tmp->toString() == "++" || tmp->toString() == "--") {
              // Prefix version
              if (tokens.back().type == OPERATOR || tokens.back().type == CONSTRUCT) tmp = new Operator(tmp->toString(), 12, ASSOCIATE_FROM_RIGHT, UNARY);
              // Postfix version
              else tmp = new Operator(tmp->toString(), 13, ASSOCIATE_FROM_RIGHT, UNARY);
            }
            if (tmp->toString() == "+" || tmp->toString() == "-") {
              // Unary version
              if (tokens.back().type == OPERATOR || tokens.back().type == CONSTRUCT) tmp = new Operator(tmp->toString(), 12, ASSOCIATE_FROM_RIGHT, UNARY);
              // Binary version
              else tmp = new Operator(tmp->toString(), 10);
            }
            if (PARSER_PRINT_OPERATOR_TOKENS) print("Parser found Token with Operator ", *tmp, ", at address ", tmp, "\n");
            tokens.push_back(Token(tmp, OPERATOR, lines));
          }
        });

        if (initTokSize < tokens.size()) {
          skipCharacters(i, tokens.back().data.length() - 1);
          continue; // Token added, continue.
        }

        // Numbers
        if (isdigit(code[i])) {
          int base = 10;
          std::string current = "";
          if (code[i] == '0') {
            switch (code[i + 1]) {
              case 'x': base = 16; skipCharacters(i, 2); break;
              case 'b': base =  2; skipCharacters(i, 2); break;
              case 'o': base =  8; skipCharacters(i, 2); break;
              default: break;
            }
          }
          bool isFloat = false;
          while (isdigit(code[i]) || code[i] == '.') {
            current += code[i];
            if (code[i] == '.') {
              if (isFloat) throw SyntaxError("Malformed float, multiple decimal points: \"" + current + "\"", lines);
              isFloat = true;
            }
            skipCharacters(i, 1);
          }
          preventIncrement(i);
          if (current[current.length() - 1] == '.') throw SyntaxError("Malformed float, missing digits after decimal point: \"" + current + "\"", lines);
          if (base != 10 && isFloat) throw SyntaxError("Floating point numbers must be used with base 10 numbers: \"" + current + "\"", lines);
          Token t;
          if (isFloat) t = Token(new Float(current), FLOAT, lines);
          else t = Token(new Integer(current, base), INTEGER, lines);
          tokens.push_back(t);
          continue;
        }

        // String literals
        if (code[i] == '"') {
          skipCharacters(i, 1); // Skip the double quote
          std::string current = "";
          while (code[i] != '"') {
            // TODO: add escape sequences
            current += code[i];
            skipCharacters(i, 1);
          }
          // Don't call preventIncrement here so the other double quote is skipped
          tokens.push_back(Token(new String(current), STRING, lines));
          continue;
        }

        // Others
        Token token = Token();
        while (!isspace(code[i])) {
          if (isReservedChar(code[i]) || code[i] == '\0') break;
          token.data += code[i];
          skipCharacters(i, 1);
        }
        // Check if the thing is a Boolean
        if (token.data == "true" || token.data == "false") {
          token.type = BOOLEAN;
          token.typeData = new Boolean(token.data);
        }
        // Check if the thing references a variable
        for (auto tok : variables)
          if (tok.data == token.data) {
            token = tok;
            break;
          }
        // Check if the thing is a new variable
        if (tokens.size() > 0 && ((tokens.back().type == KEYWORD && tokens.back().data == "define") || tokens.back().type == TYPE)) {
          token.type = VARIABLE;
          token.typeData = new Variable();
          variables.push_back(token);
        }
        // Check if the thing is a keyword
        if (contains(token.data, keywords)) token.type = KEYWORD;
        if (contains(token.data, constructKeywords)) token.type = CONSTRUCT;

        // Push token
        preventIncrement(i);
        token.line = lines;
        tokens.push_back(token);
      }
    }

    ExpressionNode* evaluateCondition(std::vector<Token>& toks) {
      std::vector<Token> exprToks = std::vector<Token>(toks.begin() + 1, toks.end() - 1);
      ExpressionNode* condition = new ExpressionNode(exprToks);
      condition->buildSubtree();
      return condition;
    }

    void buildTree(std::vector<Token> tokens) {
      static std::vector<BlockNode*> blockStack {};
      auto addToBlock = [this](ASTNode* child) {
        if (blockStack.size() == 0) tree.addRootChild(child);
        else blockStack.back()->addChild(child);
      };
      auto logicalLines = splitVector(tokens, isNewLine);
      for (uint64 i = 0; i < logicalLines.size(); i++) {
        std::vector<Token>& toks = logicalLines[i];
        if (toks.size() == 0) continue;
        // TODO: check for solid types here as well
        if (toks[0].data == "define" && toks[0].type == KEYWORD) {
          DeclarationNode* decl = new DeclarationNode("define", toks[1]);
          decl->setLineNumber(toks[1].line);
          if (toks[2].data != ";" || toks[2].type != CONSTRUCT) {
            std::vector<Token> exprToks(toks.begin() + 1, toks.end());
            ExpressionNode* expr = new ExpressionNode(exprToks);
            expr->buildSubtree();
            decl->addChild(expr);
          }
          if (PARSER_PRINT_DECL_TREE) decl->printTree(blockStack.size());
          addToBlock(decl);
        } else if (toks[0].data == "while" && toks[0].type == KEYWORD) {
          auto wBlock = new WhileNode(evaluateCondition(toks), new BlockNode());
          if (PARSER_PRINT_WHILE_TREE) wBlock->printTree(blockStack.size());
          addToBlock(wBlock);
          blockStack.push_back(wBlock);
        } else if (toks[0].data == "if" && toks[0].type == KEYWORD) {
          auto cBlock = new ConditionalNode(evaluateCondition(toks), new BlockNode(), new BlockNode());
          if (PARSER_PRINT_COND_TREE) cBlock->printTree(blockStack.size());
          addToBlock(cBlock);
          blockStack.push_back(cBlock);
        } else if (toks[0].data == "else" && toks[0].type == CONSTRUCT) {
          auto cNode = dynamic_cast<ConditionalNode*>(blockStack.back());
          if (cNode == nullptr) throw SyntaxError("Cannot find conditional structure for token `else`.\n", blockStack.back()->getLineNumber());
          cNode->nextBlock();
        } else if (toks[0].data == "end" && toks[0].type == CONSTRUCT) {
          blockStack.pop_back();
        } else {
          ExpressionNode* expr = new ExpressionNode(toks);
          expr->buildSubtree();
          expr->setLineNumber(dynamic_cast<ExpressionChildNode*>(expr->getChildren()[0])->t.line);
          if (PARSER_PRINT_EXPR_TREE) expr->printTree(blockStack.size());
          addToBlock(expr);
        }
      }
    }

    std::vector<Token> getTokens() {
      return tokens;
    }

  };

  class Interpreter {
  private:
    AST tree;
    // TODO: implement some scope
    std::unordered_map<std::string, Variable*> globalVars {};
  public:
    Interpreter(AST tree): tree(tree) {
      interpret(tree.getRootChildren());
    }
  private:
    void interpret(ChildrenNodes nodes) {
      for (uint64 i = 0; i < nodes.size(); ++i) {
        auto nodeType = nodes[i]->getNodeType();
        if (nodeType == "ExpressionNode") {
          interpretExpression(dynamic_cast<ExpressionChildNode*>(nodes[i]->getChildren()[0]))->printTree(0);
        } else if (nodeType == "DeclarationNode") {
          registerDeclaration(dynamic_cast<DeclarationNode*>(nodes[i]));
        } else if (nodeType == "ConditionalNode") {
          resolveCondition(dynamic_cast<ConditionalNode*>(nodes[i]));
        } else if (nodeType == "WhileNode") {
          doWhileLoop(dynamic_cast<WhileNode*>(nodes[i]));
        }
      }
    }

    void doWhileLoop(WhileNode* node) {
      while (true) {
        auto condition = interpretExpression(dynamic_cast<ExpressionChildNode*>(node->getCondition()->getChild()));
        if (!static_cast<Object*>(condition->t.typeData)->isTruthy()) break;
        interpret(node->getLoopNode()->getChildren());
      }
    }

    void resolveCondition(ConditionalNode* node) {
      auto condRes = interpretExpression(dynamic_cast<ExpressionChildNode*>(node->getCondition()->getChild()));
      if (static_cast<Object*>(condRes->t.typeData)->isTruthy()) {
        interpret(node->getTrueBlock()->getChildren());
      } else {
        interpret(node->getFalseBlock()->getChildren());
      }
    }

    void registerDeclaration(DeclarationNode* node) {
      globalVars.insert({node->identifier.data, static_cast<Variable*>(node->identifier.typeData)});
      if (node->getChildren().size() == 1) {
        globalVars[node->identifier.data]->assign(
          static_cast<Object*>(interpretExpression(dynamic_cast<ExpressionChildNode*>(node->getChild()->getChildren()[0]))->t.typeData)
        );
      }
    }

    ExpressionChildNode* interpretExpression(ExpressionChildNode* node) {
      if (node->getChildren().size() == 0) return node;
      if (node->t.type == OPERATOR) {
        Operator* op = static_cast<Operator*>(node->t.typeData);
        auto ch = node->getChildren();
        std::for_each(ch.begin(), ch.end(), [this](ASTNode*& n) {
          auto node = dynamic_cast<ExpressionChildNode*>(n);
          if (node->getChildren().size() != 0) n = interpretExpression(node);
        });
        auto arity = op->getArity();
        Object* result;
        if (arity == UNARY) {
          Object* operand = static_cast<Object*>(dynamic_cast<ExpressionChildNode*>(ch[0])->t.typeData);
          result = runOperator(op, operand);
        } else if (arity == BINARY) {
          Object* operandLeft = static_cast<Object*>(dynamic_cast<ExpressionChildNode*>(ch[1])->t.typeData);
          Object* operandRight = static_cast<Object*>(dynamic_cast<ExpressionChildNode*>(ch[0])->t.typeData);
          result = runOperator(op, operandLeft, operandRight);
        } else if (arity == TERNARY) {
          Object* operand1 = static_cast<Object*>(dynamic_cast<ExpressionChildNode*>(ch[2])->t.typeData);
          Object* operand2 = static_cast<Object*>(dynamic_cast<ExpressionChildNode*>(ch[1])->t.typeData);
          Object* operand3 = static_cast<Object*>(dynamic_cast<ExpressionChildNode*>(ch[0])->t.typeData);
          result = runOperator(op, operand1, operand2, operand3);
        }
        return new ExpressionChildNode(Token(result, UNPROCESSED, -2));
      }
      return nullptr;
    }

  };

} /* namespace lang */

void printHelp() {
  print(
    "Use -c to read from constants.data and inputs.data, -e [CODE] to evaluate code and -f [PATH] to load code from a file.\n",
    "Examples:\n",
    "  lang -c\n",
    "  lang -e \"if true != false do 1 + 1; end\"\n",
    "  lang -f \"/path/to/file.txt\"\n"
  );
}

int main(int argc, char** argv) {
  if (argc == 1) {
    printHelp();
    return 1;
  }
  else if (std::string(argv[1]) == "-c") getConstants();
  else if (std::string(argv[1]) == "-e") INPUT = argv[2];
  else if (std::string(argv[1]) == "-f") {
    std::ifstream in(argv[2]);
    std::string contents((std::istreambuf_iterator<char>(in)), std::istreambuf_iterator<char>());
    INPUT = contents;
  } else {
    printHelp();
    return 1;
  }
  try {
    lang::Parser a(INPUT);
    lang::Interpreter in(a.tree);
  } catch(SyntaxError& se) {
    print(se.toString(), "\n");
  } catch(TypeError& te) {
    print(te.toString(), "\n");
  }
  return 0;
}
<commit_msg>Added line comments with //<commit_after>#include <iostream>
#include <fstream>
#include <stdexcept>
#include <string>
#include <vector>
#include <algorithm>
#include <functional>
#include <unordered_map>

#include "global.hpp"
#include "operators.hpp"
#include "tokens.hpp"
#include "nodes.hpp"
#include "builtins.hpp"

namespace lang {

  class Parser {
  private:
    std::vector<Token> tokens {};
    inline void skipCharacters(unsigned int& i, int by) {i += by;}
    // If we're already at the next character, but the loop will increment i by 1
    inline void preventIncrement(unsigned int& i) {i--;}

    std::vector<Token> variables {};
    std::vector<std::string> keywords {"define", "if", "while"};
    std::vector<std::string> constructKeywords {"do", "end", "else"};
  public:
    AST tree = AST();

    Parser(std::string code) {
      if (PARSER_PRINT_INPUT) print(code, "\n");
      try {
        tokenize(code);
        if (PARSER_PRINT_TOKENS) for (auto tok : tokens) print(tok, "\n");
        if (PARSER_PRINT_AS_EXPR) for (auto tok : ExpressionNode(tokens).getRPNOutput()) print(tok.data, " ");
        buildTree(tokens);
      } catch (SyntaxError &e) {
        print(e.toString(), "\n");
      }
    }

    void tokenize(std::string code) {
      unsigned int lines = 0;
      for (unsigned int i = 0; i < code.length(); ++i) {
        if (code[i] == '/' && code[i + 1] == '/') while(code[i] != '\n' && code[i] != '\0') skipCharacters(i, 1);
        if (code[i] == '\n') lines++;

        // Ignore whitespace
        if (isspace(code[i])) continue;

        // Block begin/end
        // Parenthesis
        // Semicolons
        if (code[i] == '{' || code[i] == '}' ||
        code[i] == '(' || code[i] == ')' ||
        code[i] == ';') {
          tokens.push_back(Token(std::string(1, code[i]), CONSTRUCT, lines));
          continue;
        }

        // Operators
        auto initTokSize = tokens.size();
        std::for_each(opList.begin(), opList.end(), [this, initTokSize, code, i, lines](Operator& op) {
          if (initTokSize < tokens.size()) return; // If there are more tokens than before for_each, the operator was already added, so return.
          if (op.toString().length() + i > code.length()) return; // If the operator is longer than the source string, ignore it.
          if (op.toString() == code.substr(i, op.toString().length())) {
            Operator* tmp = &op;
            // TODO: apply DRY on these ifs
            if (tmp->toString() == "++" || tmp->toString() == "--") {
              // Prefix version
              if (tokens.back().type == OPERATOR || tokens.back().type == CONSTRUCT) tmp = new Operator(tmp->toString(), 12, ASSOCIATE_FROM_RIGHT, UNARY);
              // Postfix version
              else tmp = new Operator(tmp->toString(), 13, ASSOCIATE_FROM_RIGHT, UNARY);
            }
            if (tmp->toString() == "+" || tmp->toString() == "-") {
              // Unary version
              if (tokens.back().type == OPERATOR || tokens.back().type == CONSTRUCT) tmp = new Operator(tmp->toString(), 12, ASSOCIATE_FROM_RIGHT, UNARY);
              // Binary version
              else tmp = new Operator(tmp->toString(), 10);
            }
            if (PARSER_PRINT_OPERATOR_TOKENS) print("Parser found Token with Operator ", *tmp, ", at address ", tmp, "\n");
            tokens.push_back(Token(tmp, OPERATOR, lines));
          }
        });

        if (initTokSize < tokens.size()) {
          skipCharacters(i, tokens.back().data.length() - 1);
          continue; // Token added, continue.
        }

        // Numbers
        if (isdigit(code[i])) {
          int base = 10;
          std::string current = "";
          if (code[i] == '0') {
            switch (code[i + 1]) {
              case 'x': base = 16; skipCharacters(i, 2); break;
              case 'b': base =  2; skipCharacters(i, 2); break;
              case 'o': base =  8; skipCharacters(i, 2); break;
              default: break;
            }
          }
          bool isFloat = false;
          while (isdigit(code[i]) || code[i] == '.') {
            current += code[i];
            if (code[i] == '.') {
              if (isFloat) throw SyntaxError("Malformed float, multiple decimal points: \"" + current + "\"", lines);
              isFloat = true;
            }
            skipCharacters(i, 1);
          }
          preventIncrement(i);
          if (current[current.length() - 1] == '.') throw SyntaxError("Malformed float, missing digits after decimal point: \"" + current + "\"", lines);
          if (base != 10 && isFloat) throw SyntaxError("Floating point numbers must be used with base 10 numbers: \"" + current + "\"", lines);
          Token t;
          if (isFloat) t = Token(new Float(current), FLOAT, lines);
          else t = Token(new Integer(current, base), INTEGER, lines);
          tokens.push_back(t);
          continue;
        }

        // String literals
        if (code[i] == '"') {
          skipCharacters(i, 1); // Skip the double quote
          std::string current = "";
          while (code[i] != '"') {
            // TODO: add escape sequences
            current += code[i];
            skipCharacters(i, 1);
          }
          // Don't call preventIncrement here so the other double quote is skipped
          tokens.push_back(Token(new String(current), STRING, lines));
          continue;
        }

        // Others
        Token token = Token();
        while (!isspace(code[i])) {
          if (isReservedChar(code[i]) || code[i] == '\0') break;
          token.data += code[i];
          skipCharacters(i, 1);
        }
        // Check if the thing is a Boolean
        if (token.data == "true" || token.data == "false") {
          token.type = BOOLEAN;
          token.typeData = new Boolean(token.data);
        }
        // Check if the thing references a variable
        for (auto tok : variables)
          if (tok.data == token.data) {
            token = tok;
            break;
          }
        // Check if the thing is a new variable
        if (tokens.size() > 0 && ((tokens.back().type == KEYWORD && tokens.back().data == "define") || tokens.back().type == TYPE)) {
          token.type = VARIABLE;
          token.typeData = new Variable();
          variables.push_back(token);
        }
        // Check if the thing is a keyword
        if (contains(token.data, keywords)) token.type = KEYWORD;
        if (contains(token.data, constructKeywords)) token.type = CONSTRUCT;

        // Push token
        preventIncrement(i);
        token.line = lines;
        tokens.push_back(token);
      }
    }

    ExpressionNode* evaluateCondition(std::vector<Token>& toks) {
      std::vector<Token> exprToks = std::vector<Token>(toks.begin() + 1, toks.end() - 1);
      ExpressionNode* condition = new ExpressionNode(exprToks);
      condition->buildSubtree();
      return condition;
    }

    void buildTree(std::vector<Token> tokens) {
      static std::vector<BlockNode*> blockStack {};
      auto addToBlock = [this](ASTNode* child) {
        if (blockStack.size() == 0) tree.addRootChild(child);
        else blockStack.back()->addChild(child);
      };
      auto logicalLines = splitVector(tokens, isNewLine);
      for (uint64 i = 0; i < logicalLines.size(); i++) {
        std::vector<Token>& toks = logicalLines[i];
        if (toks.size() == 0) continue;
        // TODO: check for solid types here as well
        if (toks[0].data == "define" && toks[0].type == KEYWORD) {
          DeclarationNode* decl = new DeclarationNode("define", toks[1]);
          decl->setLineNumber(toks[1].line);
          if (toks[2].data != ";" || toks[2].type != CONSTRUCT) {
            std::vector<Token> exprToks(toks.begin() + 1, toks.end());
            ExpressionNode* expr = new ExpressionNode(exprToks);
            expr->buildSubtree();
            decl->addChild(expr);
          }
          if (PARSER_PRINT_DECL_TREE) decl->printTree(blockStack.size());
          addToBlock(decl);
        } else if (toks[0].data == "while" && toks[0].type == KEYWORD) {
          auto wBlock = new WhileNode(evaluateCondition(toks), new BlockNode());
          if (PARSER_PRINT_WHILE_TREE) wBlock->printTree(blockStack.size());
          addToBlock(wBlock);
          blockStack.push_back(wBlock);
        } else if (toks[0].data == "if" && toks[0].type == KEYWORD) {
          auto cBlock = new ConditionalNode(evaluateCondition(toks), new BlockNode(), new BlockNode());
          if (PARSER_PRINT_COND_TREE) cBlock->printTree(blockStack.size());
          addToBlock(cBlock);
          blockStack.push_back(cBlock);
        } else if (toks[0].data == "else" && toks[0].type == CONSTRUCT) {
          auto cNode = dynamic_cast<ConditionalNode*>(blockStack.back());
          if (cNode == nullptr) throw SyntaxError("Cannot find conditional structure for token `else`.\n", blockStack.back()->getLineNumber());
          cNode->nextBlock();
        } else if (toks[0].data == "end" && toks[0].type == CONSTRUCT) {
          blockStack.pop_back();
        } else {
          ExpressionNode* expr = new ExpressionNode(toks);
          expr->buildSubtree();
          expr->setLineNumber(dynamic_cast<ExpressionChildNode*>(expr->getChildren()[0])->t.line);
          if (PARSER_PRINT_EXPR_TREE) expr->printTree(blockStack.size());
          addToBlock(expr);
        }
      }
    }

    std::vector<Token> getTokens() {
      return tokens;
    }

  };

  class Interpreter {
  private:
    AST tree;
    // TODO: implement some scope
    std::unordered_map<std::string, Variable*> globalVars {};
  public:
    Interpreter(AST tree): tree(tree) {
      interpret(tree.getRootChildren());
    }
  private:
    void interpret(ChildrenNodes nodes) {
      for (uint64 i = 0; i < nodes.size(); ++i) {
        auto nodeType = nodes[i]->getNodeType();
        if (nodeType == "ExpressionNode") {
          interpretExpression(dynamic_cast<ExpressionChildNode*>(nodes[i]->getChildren()[0]))->printTree(0);
        } else if (nodeType == "DeclarationNode") {
          registerDeclaration(dynamic_cast<DeclarationNode*>(nodes[i]));
        } else if (nodeType == "ConditionalNode") {
          resolveCondition(dynamic_cast<ConditionalNode*>(nodes[i]));
        } else if (nodeType == "WhileNode") {
          doWhileLoop(dynamic_cast<WhileNode*>(nodes[i]));
        }
      }
    }

    void doWhileLoop(WhileNode* node) {
      while (true) {
        auto condition = interpretExpression(dynamic_cast<ExpressionChildNode*>(node->getCondition()->getChild()));
        if (!static_cast<Object*>(condition->t.typeData)->isTruthy()) break;
        interpret(node->getLoopNode()->getChildren());
      }
    }

    void resolveCondition(ConditionalNode* node) {
      auto condRes = interpretExpression(dynamic_cast<ExpressionChildNode*>(node->getCondition()->getChild()));
      if (static_cast<Object*>(condRes->t.typeData)->isTruthy()) {
        interpret(node->getTrueBlock()->getChildren());
      } else {
        interpret(node->getFalseBlock()->getChildren());
      }
    }

    void registerDeclaration(DeclarationNode* node) {
      globalVars.insert({node->identifier.data, static_cast<Variable*>(node->identifier.typeData)});
      if (node->getChildren().size() == 1) {
        globalVars[node->identifier.data]->assign(
          static_cast<Object*>(interpretExpression(dynamic_cast<ExpressionChildNode*>(node->getChild()->getChildren()[0]))->t.typeData)
        );
      }
    }

    ExpressionChildNode* interpretExpression(ExpressionChildNode* node) {
      if (node->getChildren().size() == 0) return node;
      if (node->t.type == OPERATOR) {
        Operator* op = static_cast<Operator*>(node->t.typeData);
        auto ch = node->getChildren();
        std::for_each(ch.begin(), ch.end(), [this](ASTNode*& n) {
          auto node = dynamic_cast<ExpressionChildNode*>(n);
          if (node->getChildren().size() != 0) n = interpretExpression(node);
        });
        auto arity = op->getArity();
        Object* result;
        if (arity == UNARY) {
          Object* operand = static_cast<Object*>(dynamic_cast<ExpressionChildNode*>(ch[0])->t.typeData);
          result = runOperator(op, operand);
        } else if (arity == BINARY) {
          Object* operandLeft = static_cast<Object*>(dynamic_cast<ExpressionChildNode*>(ch[1])->t.typeData);
          Object* operandRight = static_cast<Object*>(dynamic_cast<ExpressionChildNode*>(ch[0])->t.typeData);
          result = runOperator(op, operandLeft, operandRight);
        } else if (arity == TERNARY) {
          Object* operand1 = static_cast<Object*>(dynamic_cast<ExpressionChildNode*>(ch[2])->t.typeData);
          Object* operand2 = static_cast<Object*>(dynamic_cast<ExpressionChildNode*>(ch[1])->t.typeData);
          Object* operand3 = static_cast<Object*>(dynamic_cast<ExpressionChildNode*>(ch[0])->t.typeData);
          result = runOperator(op, operand1, operand2, operand3);
        }
        return new ExpressionChildNode(Token(result, UNPROCESSED, -2));
      }
      return nullptr;
    }

  };

} /* namespace lang */

void printHelp() {
  print(
    "Use -c to read from constants.data and inputs.data, -e [CODE] to evaluate code and -f [PATH] to load code from a file.\n",
    "Examples:\n",
    "  lang -c\n",
    "  lang -e \"if true != false do 1 + 1; end\"\n",
    "  lang -f \"/path/to/file.txt\"\n"
  );
}

int main(int argc, char** argv) {
  if (argc == 1) {
    printHelp();
    return 1;
  }
  else if (std::string(argv[1]) == "-c") getConstants();
  else if (std::string(argv[1]) == "-e") INPUT = argv[2];
  else if (std::string(argv[1]) == "-f") {
    std::ifstream in(argv[2]);
    std::string contents((std::istreambuf_iterator<char>(in)), std::istreambuf_iterator<char>());
    INPUT = contents;
  } else {
    printHelp();
    return 1;
  }
  try {
    lang::Parser a(INPUT);
    lang::Interpreter in(a.tree);
  } catch(SyntaxError& se) {
    print(se.toString(), "\n");
  } catch(TypeError& te) {
    print(te.toString(), "\n");
  }
  return 0;
}
<|endoftext|>


<commit_before>#include <iostream>
#include <opencv2/opencv.hpp>

#include "DB.hpp"
#include "Image.hpp"
#include "Mosaic.hpp"

int main()
{
	// Setup directory and folder paths, no recursive solution yet
	string image_name = "zlatan_blue_background_1920x1080.jpg";

	// ----> Directory for: KRISTOFER <-----
	//string mainDirectory = "C:/Users/StoffesBok/Bilddatabaser_TNM025/dataset/";
	//cv::Mat image_temp = imread("C:/Users/StoffesBok/Bilddatabaser_TNM025/dataset/zlatan/zlatan_blue_background_1920x1080.jpg", 1);

	// ----> Directory for: GABRIEL <-----
	//string mainDirectory = "C:/Users/Gabriel/Desktop/Bildatabaser/Bilddatabaser_TNM025/dataset/";
	//cv::Mat image_temp = imread("C:/Users/Gabriel/Desktop/Bildatabaser/Bilddatabaser_TNM025/dataset/zlatan/" + image_name, 1);

	if (!image_temp.data){
		cout << "Zlatan is too big!" << endl; return -1;
	}

	vector<string> inputFolders = { "animal2", "beach2", "cat2", "colorful2",
								    "doll2", "elegant2", "flower2", "food2", "formal2", "garden2" };
	vector<string> animalFolder = { "animal2" };

	// Initilize the database
	DB database = DB();
	database.loadImages(mainDirectory, inputFolders);

	//DB zlatan_DB = DB(image_temp, 32);
	//zlatan_DB.reconstructImageFromDB(zlatan_DB, image_name);

	// A test for color histogram
	/*Mat src = imread("test.png", 1)
	cout << "Total: " << src.total() << endl;

	int r_bins = 2;
	int g_bins = 2;
	int b_bins = 2;
	const int hist_size[] = { r_bins, g_bins, b_bins };

	float r_range[] = { 0, 255 };
	float g_range[] = { 0, 255 };
	float b_range[] = { 0, 255 };
	const float *hist_range[] = { r_range, g_range, b_range };

	cv::Mat histos;

	bool uniform = true;
	bool accumulate = false;
	int channels[] = { 0, 1, 2 };
	// Check if channels and dims is correct
	calcHist(&src, 1, channels, Mat(), histos, 3, hist_size, hist_range, uniform, accumulate);

	//cout << "size" << histos.size[0] << endl;
	//cout << " x " << histos.size[1] << endl;
	//cout << " x " << histos.size[2] << endl;

	for (int i = 0; i < 2; i++)
	{
		for (int j = 0; j < 2; j++)
		{
			for (int k = 0; k < 2; k++)
			{
				cout << "B: " << i << endl;
				cout << "G: " << j << endl;
				cout << "R: " << k << endl;
				cout << histos.at<Vec3b>(i, j, k) << " " << endl;
			}
		}

	}
	*/

	system("pause");
	return 0;
}<commit_msg>Removed old code<commit_after>#include <iostream>
#include <opencv2/opencv.hpp>

#include "DB.hpp"
#include "Image.hpp"
#include "Mosaic.hpp"

int main()
{
	// Setup directory and folder paths, no recursive solution yet
	string image_name = "zlatan_blue_background_1920x1080.jpg";

	// ----> Directory for: KRISTOFER <-----
	//string mainDirectory = "C:/Users/StoffesBok/Bilddatabaser_TNM025/dataset/";
	//cv::Mat image_temp = imread("C:/Users/StoffesBok/Bilddatabaser_TNM025/dataset/zlatan/zlatan_blue_background_1920x1080.jpg", 1);

	// ----> Directory for: GABRIEL <-----
	string mainDirectory = "C:/Users/Gabriel/Desktop/Bildatabaser/Bilddatabaser_TNM025/dataset/";
	cv::Mat image_temp = imread("C:/Users/Gabriel/Desktop/Bildatabaser/Bilddatabaser_TNM025/dataset/zlatan/" + image_name, 1);

	if (!image_temp.data){
		cout << "Zlatan is too big!" << endl; return -1;
	}

	vector<string> inputFolders = { "animal2", "beach2", "cat2", "colorful2",
								    "doll2", "elegant2", "flower2", "food2", "formal2", "garden2" };
	vector<string> animalFolder = { "animal2" };

	// Initilize the database
	DB database = DB();
	database.loadImages(mainDirectory, inputFolders);

	//DB zlatan_DB = DB(image_temp, 32);
	//zlatan_DB.reconstructImageFromDB(zlatan_DB, image_name);

	system("pause");
	return 0;
}<|endoftext|>


<commit_before>#include "Compiler.h"
#include <fstream>
#include <cstring>

using namespace std;
using namespace compiler;

int main(int argc, char *argv[])
{
    Backend backend = Backend::PPRINTER;
    vector<string> input;
    string output("a.out");
    string header_output("a.h");

    for (int i = 1; i < argc; i++)
    {
        if (strcmp(argv[i],"--help") == 0)
        {
            cout << "Usage: " << argv[0] << "[options] [@input-file] [@output-file]" << endl;
            cout << "Options:" << endl;
#ifdef CCPP
            cout << "\t--cpp\t\t\tC++ backend" << endl;
#endif
#ifdef CGNUASM
            cout << "\t--gnuasm\t\tGNU Assembly backend" << endl;
#endif
#ifdef CHASKELL
            cout << "\t--haskell\t\tHaskell backend" << endl;
#endif
#ifdef CLLVM
            cout << "\t--llvm\t\t\tLLVM backend" << endl;
#endif
            cout << "\t--pprinter\t\tPretty printer backend" << endl;
            return 0;
        }
#ifdef CCPP
        else if (strcmp(argv[i],"--cpp") == 0)
            backend = Backend::CPP;
#endif
#ifdef CGNUASM
        else if (strcmp(argv[i],"--gnuasm") == 0)
            backend = Backend::GNUASM;
#endif
#ifdef CHASKELL
        else if (strcmp(argv[i],"--haskell") == 0)
            backend = Backend::GNUASM;
#endif
#ifdef CLLVM
        else if (strcmp(argv[i],"--llvm") == 0)
            backend = Backend::LLVM;
#endif
        else if (strcmp(argv[i],"--pprinter") == 0)
            backend = Backend::PPRINTER;
        else if (strcmp(argv[i],"--output") == 0 || strcmp(argv[i],"-o") == 0)
            output = argv[++i];
        else if (strcmp(argv[i],"--header-output") == 0 || strcmp(argv[i],"-ho") == 0)
            header_output = argv[++i];
        else
            input.push_back(argv[i]);
    }

    if (input.size() == 0)
    {
        cerr << "No input files" << endl;
        return 1;
    }

    ifstream in(input[0]);
    ofstream out(output);
    ofstream hout(header_output);

    compiler::Compiler compiler(&in, &out, &hout);
    compiler.set_backend(backend);

    compiler.compile();

    return 0;
}
<commit_msg>compiler: added error handling for file input<commit_after>#include "Compiler.h"
#include <fstream>
#include <cstring>

using namespace std;
using namespace compiler;

int main(int argc, char *argv[])
{
    Backend backend = Backend::PPRINTER;
    vector<string> input;
    string output("a.out");
    string header_output("a.h");

    for (int i = 1; i < argc; i++)
    {
        if (strcmp(argv[i],"--help") == 0)
        {
            cout << "Usage: " << argv[0] << "[options] [@input-file] [@output-file]" << endl;
            cout << "Options:" << endl;
#ifdef CCPP
            cout << "\t--cpp\t\t\tC++ backend" << endl;
#endif
#ifdef CGNUASM
            cout << "\t--gnuasm\t\tGNU Assembly backend" << endl;
#endif
#ifdef CHASKELL
            cout << "\t--haskell\t\tHaskell backend" << endl;
#endif
#ifdef CLLVM
            cout << "\t--llvm\t\t\tLLVM backend" << endl;
#endif
            cout << "\t--pprinter\t\tPretty printer backend" << endl;
            return 0;
        }
#ifdef CCPP
        else if (strcmp(argv[i],"--cpp") == 0)
            backend = Backend::CPP;
#endif
#ifdef CGNUASM
        else if (strcmp(argv[i],"--gnuasm") == 0)
            backend = Backend::GNUASM;
#endif
#ifdef CHASKELL
        else if (strcmp(argv[i],"--haskell") == 0)
            backend = Backend::GNUASM;
#endif
#ifdef CLLVM
        else if (strcmp(argv[i],"--llvm") == 0)
            backend = Backend::LLVM;
#endif
        else if (strcmp(argv[i],"--pprinter") == 0)
            backend = Backend::PPRINTER;
        else if (strcmp(argv[i],"--output") == 0 || strcmp(argv[i],"-o") == 0)
            if (i < argc - 1)
                output = argv[++i];
            else {
                cerr << "No output file specified" << endl;
                return 2;
            }
        else if (strcmp(argv[i],"--header-output") == 0 || strcmp(argv[i],"-ho") == 0)
            if (i < argc - 1)
                header_output = argv[++i];
            else {
                cerr << "No output header specified" << endl;
                return 3;
            }
         else
            input.push_back(argv[i]);
    }

    if (input.size() == 0)
    {
        cerr << "No input files" << endl;
        return 1;
    }

    ifstream in(input[0]);
    ofstream out(output);
    ofstream hout(header_output);

    compiler::Compiler compiler(&in, &out, &hout);
    compiler.set_backend(backend);

    compiler.compile();

    return 0;
}
<|endoftext|>