<commit_before>#include "NxFileSystem.hpp"
#include "../MemoryFile.hpp"
#include "../InputStream.hpp"
#include "../OutputStream.hpp"
#include "../Exception.hpp"
#include <nn/fs.h>

BEGIN_INANITY_PLATFORM

/// Check result of operation.
/** If message is null, simply return result. Otherwise throw an exception. */
static bool NxCheckResult(const nn::Result& result, const char* message = nullptr)
{
	if(result.IsSuccess()) return true;

	if(!message) return false;

	ptr<Exception> exception;
	if(nn::fs::ResultPathNotFound::Includes(result))
	{
		exception = NEW(Exception("Path not found"));
	}
	else if(nn::fs::ResultPathAlreadyExists::Includes(result))
	{
		exception = NEW(Exception("Path already exists"));
	}
	else if(nn::fs::ResultTargetLocked::Includes(result))
	{
		exception = NEW(Exception("Target locked"));
	}
	else if(nn::fs::ResultUsableSpaceNotEnough::Includes(result))
	{
		exception = NEW(Exception("Usable space not enough"));
	}

	if(exception)
		THROW_SECONDARY(message, exception);
	else
		THROW(message);
}

class NxFileHandle
{
private:
	nn::fs::FileHandle fileHandle = { nullptr };

public:
	NxFileHandle() {}

	NxFileHandle(nn::fs::FileHandle fileHandle)
	: fileHandle(fileHandle) {}

	NxFileHandle(const NxFileHandle&) = delete;

	NxFileHandle(NxFileHandle&& fileHandle)
	: fileHandle(fileHandle.fileHandle)
	{
		fileHandle.fileHandle = { nullptr };
	}

	~NxFileHandle()
	{
		if(fileHandle.handle)
			nn::fs::CloseFile(fileHandle);
	}

	operator nn::fs::FileHandle() const
	{
		return fileHandle;
	}

	nn::fs::FileHandle* operator&()
	{
		return &fileHandle;
	}
};

class NxDirectoryHandle
{
private:
	nn::fs::DirectoryHandle directoryHandle = { nullptr };

public:
	NxDirectoryHandle() {}

	NxDirectoryHandle(nn::fs::DirectoryHandle directoryHandle)
	: directoryHandle(directoryHandle) {}

	NxDirectoryHandle(const NxDirectoryHandle&) = delete;

	NxDirectoryHandle(NxDirectoryHandle&& directoryHandle)
	: directoryHandle(directoryHandle.directoryHandle)
	{
		directoryHandle.directoryHandle = { nullptr };
	}

	~NxDirectoryHandle()
	{
		if(directoryHandle.handle)
			nn::fs::CloseDirectory(directoryHandle);
	}

	operator nn::fs::DirectoryHandle() const
	{
		return directoryHandle;
	}

	nn::fs::DirectoryHandle* operator&()
	{
		return &directoryHandle;
	}
};

class NxInputStream : public InputStream
{
private:
	NxFileHandle fileHandle;
	int64_t offset = 0;
	int64_t fileSize = 0;

public:
	NxInputStream(NxFileHandle&& fileHandle, int64_t fileSize)
	: fileHandle(std::move(fileHandle)), fileSize(fileSize) {}

	size_t Read(void* data, size_t size) override
	{
		int64_t end = offset + (int64_t)size;
		if(end > fileSize) end = fileSize;
		size_t read = end - offset;
		NxCheckResult(nn::fs::ReadFile(fileHandle, offset, data, read), "Can't read from Nx input stream");
		return read;
	}

	bigsize_t Skip(bigsize_t size) override
	{
		int64_t newOffset = offset + (int64_t)size;
		if(newOffset > fileSize) newOffset = fileSize;
		bigsize_t skipped = newOffset - offset;
		offset = newOffset;
		return skipped;
	}

	bool IsAtEnd() const override
	{
		return offset >= fileSize;
	}
};

class NxOutputStream : public OutputStream
{
	NxFileHandle fileHandle;
	int64_t offset = 0;

public:
	NxOutputStream(NxFileHandle&& fileHandle)
	: fileHandle(std::move(fileHandle)) {}

	void Write(const void* data, size_t size) override
	{
		NxCheckResult(nn::fs::WriteFile(fileHandle, offset, data, size, nn::fs::WriteOption::MakeValue(0)), "Can't write to Nx output stream");
		offset += size;
	}
};

NxFileSystem::NxFileSystem(const String& userFolderName)
: folderName(userFolderName)
{
	if(!folderName.empty() && folderName[folderName.length() - 1] != '/')
		folderName += "/";
}

String NxFileSystem::GetFullName(const String& fileName) const
{
	if(fileName.empty() || fileName[0] != '/')
		THROW("File name should begin with slash");

	return folderName + fileName.substr(1);
}

ptr<File> NxFileSystem::LoadFile(const String& fileName)
{
	// open file
	NxFileHandle fileHandle;
	NxCheckResult(nn::fs::OpenFile(&fileHandle, GetFullName(fileName).c_str(), nn::fs::OpenMode_Read), "Can't open file");

	// get size
	int64_t fileSize;
	NxCheckResult(nn::fs::GetFileSize(&fileSize, fileHandle), "Can't get file size");

	// load data
	ptr<File> file = NEW(MemoryFile((size_t)fileSize));
	NxCheckResult(nn::fs::ReadFile(fileHandle, 0, file->GetData(), file->GetSize()), "Can't read file");

	return file;
}

ptr<File> NxFileSystem::TryLoadFile(const String& fileName)
{
	// open file
	NxFileHandle fileHandle;
	if(!NxCheckResult(nn::fs::OpenFile(&fileHandle, GetFullName(fileName).c_str(), nn::fs::OpenMode_Read))) return nullptr;

	// get size
	int64_t fileSize;
	if(!NxCheckResult(nn::fs::GetFileSize(&fileSize, fileHandle))) return nullptr;

	// load data
	ptr<File> file = NEW(MemoryFile((size_t)fileSize));
	if(!NxCheckResult(nn::fs::ReadFile(fileHandle, 0, file->GetData(), file->GetSize()))) return nullptr;

	return file;
}

ptr<InputStream> NxFileSystem::LoadStream(const String& fileName)
{
	// open file
	NxFileHandle fileHandle;
	NxCheckResult(nn::fs::OpenFile(&fileHandle, GetFullName(fileName).c_str(), nn::fs::OpenMode_Read), "Can't open file");

	// get size
	int64_t fileSize;
	NxCheckResult(nn::fs::GetFileSize(&fileSize, fileHandle), "Can't get file size");

	// create stream
	return NEW(NxInputStream(std::move(fileHandle), fileSize));
}

void NxFileSystem::SaveFile(ptr<File> file, const String& fileName)
{
	// open file
	NxFileHandle fileHandle;
	NxCheckResult(nn::fs::OpenFile(&fileHandle, GetFullName(fileName).c_str(), nn::fs::OpenMode_Write | nn::fs::OpenMode_AllowAppend), "Can't open file");

	// write data
	NxCheckResult(nn::fs::WriteFile(fileHandle, 0, file->GetData(), file->GetSize(), nn::fs::WriteOption::MakeValue(0)), "Can't write file");
}

ptr<OutputStream> NxFileSystem::SaveStream(const String& fileName)
{
	// open file
	NxFileHandle fileHandle;
	NxCheckResult(nn::fs::OpenFile(&fileHandle, GetFullName(fileName).c_str(), nn::fs::OpenMode_Write | nn::fs::OpenMode_AllowAppend), "Can't open file");

	// create stream
	return NEW(NxOutputStream(std::move(fileHandle)));
}

#if defined(_DEBUG)

time_t NxFileSystem::GetFileMTime(const String& fileName)
{
	nn::fs::FileTimeStamp fileTimeStamp;
	NxCheckResult(nn::fs::GetFileTimeStampForDebug(&fileTimeStamp, GetFullName(fileName).c_str()), "Can't get file timestamp");
	return fileTimeStamp.modify.value;
}

#endif

void NxFileSystem::GetDirectoryEntries(const String& directoryName, std::vector<String>& entries) const
{
	// open directory
	NxDirectoryHandle directoryHandle;
	NxCheckResult(nn::fs::OpenDirectory(&directoryHandle, GetFullName(directoryName).c_str(), nn::fs::OpenDirectoryMode_All), "Can't open directory");

	// get entry count
	// we get entries in batches anyway, but we don't want to allocate too much for small directories
	int64_t entriesBufferSize;
	NxCheckResult(nn::fs::GetDirectoryEntryCount(&entriesBufferSize, directoryHandle), "Can't get directory entry count");
	if(entriesBufferSize <= 0) return;
	if(entriesBufferSize > 64) entriesBufferSize = 64;

	nn::fs::DirectoryEntry tempEntries[entriesBufferSize]; // VLA - GNU C++ extension
	for(;;)
	{
		int64_t entriesStored;
		NxCheckResult(nn::fs::ReadDirectory(&entriesStored, tempEntries, directoryHandle, entriesBufferSize), "Can't get directory entries");

		if(entriesStored <= 0) break;

		for(int64_t i = 0; i < entriesStored; ++i)
			entries.push_back(tempEntries[i].name);
	}
}

void NxFileSystem::MakeDirectory(const String& directoryName)
{
	NxCheckResult(nn::fs::CreateDirectory(GetFullName(directoryName).c_str()), "Can't create directory");
}

void NxFileSystem::DeleteEntry(const String& entryName)
{
	// support only deleting files for now
	NxCheckResult(nn::fs::DeleteFile(GetFullName(entryName).c_str()), "Can't delete file");
}

FileSystem::EntryType NxFileSystem::GetEntryType(const String& entryName) const
{
	nn::fs::DirectoryEntryType entryType;
	NxCheckResult(nn::fs::GetEntryType(&entryType, GetFullName(entryName).c_str()), "Can't get entry type");
	switch(entryType)
	{
	case nn::fs::DirectoryEntryType_File:
		return entryTypeFile;
	case nn::fs::DirectoryEntryType_Directory:
		return entryTypeDirectory;
	default:
		return entryTypeOther;
	}
}

ptr<FileSystem> NxFileSystem::GetSubFileSystem(const String& folderName)
{
	return NEW(NxFileSystem(GetFullName(folderName)));
}

END_INANITY_PLATFORM
<commit_msg>fix writing to files in NxFileSystem<commit_after>#include "NxFileSystem.hpp"
#include "../MemoryFile.hpp"
#include "../InputStream.hpp"
#include "../OutputStream.hpp"
#include "../Exception.hpp"
#include <nn/fs.h>

BEGIN_INANITY_PLATFORM

/// Check result of operation.
/** If message is null, simply return result. Otherwise throw an exception. */
static bool NxCheckResult(const nn::Result& result, const char* message = nullptr)
{
	if(result.IsSuccess()) return true;

	if(!message) return false;

	ptr<Exception> exception;
	if(nn::fs::ResultPathNotFound::Includes(result))
	{
		exception = NEW(Exception("Path not found"));
	}
	else if(nn::fs::ResultPathAlreadyExists::Includes(result))
	{
		exception = NEW(Exception("Path already exists"));
	}
	else if(nn::fs::ResultTargetLocked::Includes(result))
	{
		exception = NEW(Exception("Target locked"));
	}
	else if(nn::fs::ResultUsableSpaceNotEnough::Includes(result))
	{
		exception = NEW(Exception("Usable space not enough"));
	}

	if(exception)
		THROW_SECONDARY(message, exception);
	else
		THROW(message);
}

class NxFileHandle
{
private:
	nn::fs::FileHandle fileHandle = { nullptr };

public:
	NxFileHandle() {}

	NxFileHandle(nn::fs::FileHandle fileHandle)
	: fileHandle(fileHandle) {}

	NxFileHandle(const NxFileHandle&) = delete;

	NxFileHandle(NxFileHandle&& fileHandle)
	: fileHandle(fileHandle.fileHandle)
	{
		fileHandle.fileHandle = { nullptr };
	}

	~NxFileHandle()
	{
		if(fileHandle.handle)
			nn::fs::CloseFile(fileHandle);
	}

	operator nn::fs::FileHandle() const
	{
		return fileHandle;
	}

	nn::fs::FileHandle* operator&()
	{
		return &fileHandle;
	}
};

class NxDirectoryHandle
{
private:
	nn::fs::DirectoryHandle directoryHandle = { nullptr };

public:
	NxDirectoryHandle() {}

	NxDirectoryHandle(nn::fs::DirectoryHandle directoryHandle)
	: directoryHandle(directoryHandle) {}

	NxDirectoryHandle(const NxDirectoryHandle&) = delete;

	NxDirectoryHandle(NxDirectoryHandle&& directoryHandle)
	: directoryHandle(directoryHandle.directoryHandle)
	{
		directoryHandle.directoryHandle = { nullptr };
	}

	~NxDirectoryHandle()
	{
		if(directoryHandle.handle)
			nn::fs::CloseDirectory(directoryHandle);
	}

	operator nn::fs::DirectoryHandle() const
	{
		return directoryHandle;
	}

	nn::fs::DirectoryHandle* operator&()
	{
		return &directoryHandle;
	}
};

class NxInputStream : public InputStream
{
private:
	NxFileHandle fileHandle;
	int64_t offset = 0;
	int64_t fileSize = 0;

public:
	NxInputStream(NxFileHandle&& fileHandle, int64_t fileSize)
	: fileHandle(std::move(fileHandle)), fileSize(fileSize) {}

	size_t Read(void* data, size_t size) override
	{
		int64_t end = offset + (int64_t)size;
		if(end > fileSize) end = fileSize;
		size_t read = end - offset;
		NxCheckResult(nn::fs::ReadFile(fileHandle, offset, data, read), "Can't read from Nx input stream");
		return read;
	}

	bigsize_t Skip(bigsize_t size) override
	{
		int64_t newOffset = offset + (int64_t)size;
		if(newOffset > fileSize) newOffset = fileSize;
		bigsize_t skipped = newOffset - offset;
		offset = newOffset;
		return skipped;
	}

	bool IsAtEnd() const override
	{
		return offset >= fileSize;
	}
};

class NxOutputStream : public OutputStream
{
	NxFileHandle fileHandle;
	int64_t offset = 0;

public:
	NxOutputStream(NxFileHandle&& fileHandle)
	: fileHandle(std::move(fileHandle)) {}

	~NxOutputStream()
	{
		// flush before close
		nn::fs::FlushFile(fileHandle);
	}

	void Write(const void* data, size_t size) override
	{
		NxCheckResult(nn::fs::WriteFile(fileHandle, offset, data, size, nn::fs::WriteOption::MakeValue(0)), "Can't write to Nx output stream");
		offset += size;
	}
};

NxFileSystem::NxFileSystem(const String& userFolderName)
: folderName(userFolderName)
{
	if(!folderName.empty() && folderName[folderName.length() - 1] != '/')
		folderName += "/";
}

String NxFileSystem::GetFullName(const String& fileName) const
{
	if(fileName.empty() || fileName[0] != '/')
		THROW("File name should begin with slash");

	return folderName + fileName.substr(1);
}

ptr<File> NxFileSystem::LoadFile(const String& fileName)
{
	// open file
	NxFileHandle fileHandle;
	NxCheckResult(nn::fs::OpenFile(&fileHandle, GetFullName(fileName).c_str(), nn::fs::OpenMode_Read), "Can't open file");

	// get size
	int64_t fileSize;
	NxCheckResult(nn::fs::GetFileSize(&fileSize, fileHandle), "Can't get file size");

	// load data
	ptr<File> file = NEW(MemoryFile((size_t)fileSize));
	NxCheckResult(nn::fs::ReadFile(fileHandle, 0, file->GetData(), file->GetSize()), "Can't read file");

	return file;
}

ptr<File> NxFileSystem::TryLoadFile(const String& fileName)
{
	// open file
	NxFileHandle fileHandle;
	if(!NxCheckResult(nn::fs::OpenFile(&fileHandle, GetFullName(fileName).c_str(), nn::fs::OpenMode_Read))) return nullptr;

	// get size
	int64_t fileSize;
	if(!NxCheckResult(nn::fs::GetFileSize(&fileSize, fileHandle))) return nullptr;

	// load data
	ptr<File> file = NEW(MemoryFile((size_t)fileSize));
	if(!NxCheckResult(nn::fs::ReadFile(fileHandle, 0, file->GetData(), file->GetSize()))) return nullptr;

	return file;
}

ptr<InputStream> NxFileSystem::LoadStream(const String& fileName)
{
	// open file
	NxFileHandle fileHandle;
	NxCheckResult(nn::fs::OpenFile(&fileHandle, GetFullName(fileName).c_str(), nn::fs::OpenMode_Read), "Can't open file");

	// get size
	int64_t fileSize;
	NxCheckResult(nn::fs::GetFileSize(&fileSize, fileHandle), "Can't get file size");

	// create stream
	return NEW(NxInputStream(std::move(fileHandle), fileSize));
}

void NxFileSystem::SaveFile(ptr<File> file, const String& fileName)
{
	String fullName = GetFullName(fileName);

	// create file if not exist
	nn::fs::CreateFile(fullName.c_str(), 0);

	// open file
	NxFileHandle fileHandle;
	NxCheckResult(nn::fs::OpenFile(&fileHandle, fullName.c_str(), nn::fs::OpenMode_Write | nn::fs::OpenMode_AllowAppend), "Can't open file");

	// write data
	NxCheckResult(nn::fs::WriteFile(fileHandle, 0, file->GetData(), file->GetSize(), nn::fs::WriteOption::MakeValue(0)), "Can't write file");

	// apparently flushing after writing is mandatory (otherwise it crashes on CloseFile)
	NxCheckResult(nn::fs::FlushFile(fileHandle), "Can't flush file");;
}

ptr<OutputStream> NxFileSystem::SaveStream(const String& fileName)
{
	// open file
	NxFileHandle fileHandle;
	NxCheckResult(nn::fs::OpenFile(&fileHandle, GetFullName(fileName).c_str(), nn::fs::OpenMode_Write | nn::fs::OpenMode_AllowAppend), "Can't open file");

	// create stream
	return NEW(NxOutputStream(std::move(fileHandle)));
}

#if defined(_DEBUG)

time_t NxFileSystem::GetFileMTime(const String& fileName)
{
	nn::fs::FileTimeStamp fileTimeStamp;
	NxCheckResult(nn::fs::GetFileTimeStampForDebug(&fileTimeStamp, GetFullName(fileName).c_str()), "Can't get file timestamp");
	return fileTimeStamp.modify.value;
}

#endif

void NxFileSystem::GetDirectoryEntries(const String& directoryName, std::vector<String>& entries) const
{
	// open directory
	NxDirectoryHandle directoryHandle;
	NxCheckResult(nn::fs::OpenDirectory(&directoryHandle, GetFullName(directoryName).c_str(), nn::fs::OpenDirectoryMode_All), "Can't open directory");

	// get entry count
	// we get entries in batches anyway, but we don't want to allocate too much for small directories
	int64_t entriesBufferSize;
	NxCheckResult(nn::fs::GetDirectoryEntryCount(&entriesBufferSize, directoryHandle), "Can't get directory entry count");
	if(entriesBufferSize <= 0) return;
	if(entriesBufferSize > 64) entriesBufferSize = 64;

	nn::fs::DirectoryEntry tempEntries[entriesBufferSize]; // VLA - GNU C++ extension
	for(;;)
	{
		int64_t entriesStored;
		NxCheckResult(nn::fs::ReadDirectory(&entriesStored, tempEntries, directoryHandle, entriesBufferSize), "Can't get directory entries");

		if(entriesStored <= 0) break;

		for(int64_t i = 0; i < entriesStored; ++i)
			entries.push_back(tempEntries[i].name);
	}
}

void NxFileSystem::MakeDirectory(const String& directoryName)
{
	NxCheckResult(nn::fs::CreateDirectory(GetFullName(directoryName).c_str()), "Can't create directory");
}

void NxFileSystem::DeleteEntry(const String& entryName)
{
	// support only deleting files for now
	NxCheckResult(nn::fs::DeleteFile(GetFullName(entryName).c_str()), "Can't delete file");
}

FileSystem::EntryType NxFileSystem::GetEntryType(const String& entryName) const
{
	nn::fs::DirectoryEntryType entryType;
	NxCheckResult(nn::fs::GetEntryType(&entryType, GetFullName(entryName).c_str()), "Can't get entry type");
	switch(entryType)
	{
	case nn::fs::DirectoryEntryType_File:
		return entryTypeFile;
	case nn::fs::DirectoryEntryType_Directory:
		return entryTypeDirectory;
	default:
		return entryTypeOther;
	}
}

ptr<FileSystem> NxFileSystem::GetSubFileSystem(const String& folderName)
{
	return NEW(NxFileSystem(GetFullName(folderName)));
}

END_INANITY_PLATFORM
<|endoftext|>


<commit_before>#include <phypp.hpp>

void print_help();

int phypp_main(int argc, char* argv[]) {
    if (argc < 3) {
        print_help();
        return 0;
    }

    vec1s suffix;
    std::string fitmask;
    std::string outdir;
    bool residuals = false;
    read_args(argc-2, argv+2, arg_list(suffix, fitmask, residuals, outdir));

    if (!outdir.empty()) {
        outdir = file::directorize(outdir);
        file::mkdir(outdir);
    }

    // Read cube
    fits::input_image fimg(argv[1]);
    vec3d flx;
    fimg.reach_hdu(1);
    fimg.read(flx);
    // Read uncertainty
    vec3d err;
    fimg.reach_hdu(2);
    fimg.read(err);

    // Get wavelength WCS
    double crpix, crval, cdelt;
    fimg.read_keyword("CRPIX3", crpix);
    fimg.read_keyword("CRVAL3", crval);
    fimg.read_keyword("CDELT3", cdelt);

    // Read fit mask (optional)
    vec2d mask; {
        if (!fitmask.empty()) {
            fits::read(fitmask, mask);
        } else {
            mask = replicate(1, flx.dims[1], flx.dims[2]);
        }
    }

    // Read models
    vec2d models;
    uint_t nmodel; {
        fits::input_image fmodel(argv[2]);
        if (fmodel.axis_count() == 2) {
            nmodel = 1;

        }
        vec3d tmp;
        fits::read(argv[2], tmp);

        nmodel = tmp.dims[0];
        models = reform(tmp, nmodel, tmp.dims[1]*tmp.dims[2]);

        // Make sure each model has unit integral
        for (uint_t i : range(nmodel)) {
            vec1u idg = where(is_finite(models(i,_)));
            models(i,_) /= total(models(i,_)[idg]);
        }

        if (suffix.size() != nmodel) {
            error("please provide as many suffixes as there are models (", nmodel,
                ") in suffix=[...]");
            return 1;
        }
    }

    // Adjust suffixes to include "_"
    if (nmodel == 1 && suffix.empty()) {
        suffix = {""};
    }

    for (std::string& s : suffix) {
        if (!s.empty() && s[0] != '_') {
            s = "_"+s;
        }
    }

    // Do the fit wavelength by wavelength
    vec2d flx1d(flx.dims[0], nmodel);
    vec2d err1d(flx.dims[0], nmodel);
    for (uint_t l : range(flx.dims[0])) {
        vec1u idg = where(is_finite(flx(l,_,_)) && is_finite(err(l,_,_)) && mask > 0.0);
        auto res = linfit_pack(flx(l,_,_)[idg], err(l,_,_)[idg], models(_,idg));
        flx1d(l,_) = res.params;
        err1d(l,_) = res.errors;
    }

    // Write spectra to disk
    std::string filebase = outdir+file::get_basename(erase_end(argv[1], ".fits"));
    for (uint_t i : range(nmodel)) {
        fits::output_image fspec(filebase+suffix[i]+"_spec.fits");
        fspec.write(vec1d(0)); // empty primary extension, KMOS convention
        fspec.reach_hdu(1);
        fspec.write(flx1d(_,i));
        fspec.write_keyword("CRPIX1", crpix);
        fspec.write_keyword("CRVAL1", crval);
        fspec.write_keyword("CDELT1", cdelt);
        fspec.reach_hdu(2);
        fspec.write(err1d(_,i));
        fspec.write_keyword("CRPIX1", crpix);
        fspec.write_keyword("CRVAL1", crval);
        fspec.write_keyword("CDELT1", cdelt);
    }

    // Compute residuals if asked
    if (residuals) {
        fimg.reach_hdu(1);
        fits::header hdr = fimg.read_header();

        for (uint_t i : range(nmodel)) {
            vec3d res = flx;
            for (uint_t l : range(res.dims[0])) {
                flatten(res(l,_,_)) -= flx1d(l,i)*models(i,_);
            }

            fits::output_image fspec(filebase+suffix[i]+"_residual.fits");
            fspec.write(vec3d(0,0,0)); // empty primary HDU
            fspec.reach_hdu(1);
            fspec.write(res);
            fspec.write_header(hdr);
            fspec.reach_hdu(2);
            fspec.write(err);
            fspec.write_header(hdr);
        }
    }

    return 0;
}

void print_help() {
    using namespace format;

    print("multispecfit v1.0");
    print("usage: multispecfit <kmos_cube.fits> <models.fits> suffix=[...] [options]");
    print("");
    print("Main parameters:");
    paragraph("'kmos_cube.fits' must be a cube created by the KMOS pipeline, with 3 "
        "extensions: the first is empty (KMOS convention), the second contains the flux, "
        "and the third contains the uncertainty. 'models.fits' must be a cube created by "
        "yourself (e.g., with IDL or Python) containing the models to fit. Each slice of "
        "this cube corresponds to a different model. Note that you have to include a "
        "uniform model (all pixels equal to 1) if you want to fit a constant background. "
        "Finally, 'suffix' must contain the names of each model, with the following "
        "format: [name1,name1,name3,...]. There must be as many suffixes as there are "
        "models in the 'models.fits' file.");
    print("Available options:");
    bullet("fitmask=...", "Must be a 2D FITS file containing the mask to define the fitting "
        "region. Only the pixels with a non-zero value in the mask will be used "
        "(default: all valid pixels are used).");
    bullet("residuals", "Set this flag if you want the program to generate residual cubes, "
        "removing the contribution of each component separately (default: do not "
        "generate residuals)");
    bullet("outdir", "Name of the directory into which the output files should be created. "
        "Default is the current directory.");
}
<commit_msg>multispecfit can now fit a single model as a 2D file rather than cube<commit_after>#include <phypp.hpp>

void print_help();

int phypp_main(int argc, char* argv[]) {
    if (argc < 3) {
        print_help();
        return 0;
    }

    vec1s suffix;
    std::string fitmask;
    std::string outdir;
    bool residuals = false;
    read_args(argc-2, argv+2, arg_list(suffix, fitmask, residuals, outdir));

    if (!outdir.empty()) {
        outdir = file::directorize(outdir);
        file::mkdir(outdir);
    }

    // Read cube
    fits::input_image fimg(argv[1]);
    vec3d flx;
    fimg.reach_hdu(1);
    fimg.read(flx);
    // Read uncertainty
    vec3d err;
    fimg.reach_hdu(2);
    fimg.read(err);

    // Get wavelength WCS
    double crpix, crval, cdelt;
    fimg.read_keyword("CRPIX3", crpix);
    fimg.read_keyword("CRVAL3", crval);
    fimg.read_keyword("CDELT3", cdelt);

    // Read fit mask (optional)
    vec2d mask; {
        if (!fitmask.empty()) {
            fits::read(fitmask, mask);
        } else {
            mask = replicate(1, flx.dims[1], flx.dims[2]);
        }
    }

    // Read models
    vec2d models;
    uint_t nmodel; {
        fits::input_image fmodel(argv[2]);
        if (fmodel.axis_count() == 2) {
            nmodel = 1;
            fits::read(argv[2], models);
            models = reform(models, 1, models.size());
        } else {
            vec3d tmp;
            fits::read(argv[2], tmp);

            nmodel = tmp.dims[0];
            models = reform(tmp, nmodel, tmp.dims[1]*tmp.dims[2]);
        }

        // Make sure each model has unit integral
        for (uint_t i : range(nmodel)) {
            vec1u idg = where(is_finite(models(i,_)));
            models(i,_) /= total(models(i,_)[idg]);
        }

        if (suffix.size() != nmodel && nmodel != 1) {
            error("please provide as many suffixes as there are models (", nmodel,
                ") in suffix=[...]");
            return 1;
        }
    }

    // Adjust suffixes to include "_"
    if (nmodel == 1 && suffix.empty()) {
        suffix = {""};
    }

    for (std::string& s : suffix) {
        if (!s.empty() && s[0] != '_') {
            s = "_"+s;
        }
    }

    // Do the fit wavelength by wavelength
    vec2d flx1d(flx.dims[0], nmodel);
    vec2d err1d(flx.dims[0], nmodel);
    for (uint_t l : range(flx.dims[0])) {
        vec1u idg = where(is_finite(flx(l,_,_)) && is_finite(err(l,_,_)) && mask > 0.0);
        auto res = linfit_pack(flx(l,_,_)[idg], err(l,_,_)[idg], models(_,idg));
        flx1d(l,_) = res.params;
        err1d(l,_) = res.errors;
    }

    // Write spectra to disk
    std::string filebase = outdir+file::get_basename(erase_end(argv[1], ".fits"));
    for (uint_t i : range(nmodel)) {
        fits::output_image fspec(filebase+suffix[i]+"_spec.fits");
        fspec.write(vec1d(0)); // empty primary extension, KMOS convention
        fspec.reach_hdu(1);
        fspec.write(flx1d(_,i));
        fspec.write_keyword("CRPIX1", crpix);
        fspec.write_keyword("CRVAL1", crval);
        fspec.write_keyword("CDELT1", cdelt);
        fspec.reach_hdu(2);
        fspec.write(err1d(_,i));
        fspec.write_keyword("CRPIX1", crpix);
        fspec.write_keyword("CRVAL1", crval);
        fspec.write_keyword("CDELT1", cdelt);
    }

    // Compute residuals if asked
    if (residuals) {
        fimg.reach_hdu(1);
        fits::header hdr = fimg.read_header();

        for (uint_t i : range(nmodel)) {
            vec3d res = flx;
            for (uint_t l : range(res.dims[0])) {
                flatten(res(l,_,_)) -= flx1d(l,i)*models(i,_);
            }

            fits::output_image fspec(filebase+suffix[i]+"_residual.fits");
            fspec.write(vec3d(0,0,0)); // empty primary HDU
            fspec.reach_hdu(1);
            fspec.write(res);
            fspec.write_header(hdr);
            fspec.reach_hdu(2);
            fspec.write(err);
            fspec.write_header(hdr);
        }
    }

    return 0;
}

void print_help() {
    using namespace format;

    print("multispecfit v1.0");
    print("usage: multispecfit <kmos_cube.fits> <models.fits> suffix=[...] [options]");
    print("");
    print("Main parameters:");
    paragraph("'kmos_cube.fits' must be a cube created by the KMOS pipeline, with 3 "
        "extensions: the first is empty (KMOS convention), the second contains the flux, "
        "and the third contains the uncertainty. 'models.fits' must be a cube created by "
        "yourself (e.g., with IDL or Python) containing the models to fit. Each slice of "
        "this cube corresponds to a different model. Note that you have to include a "
        "uniform model (all pixels equal to 1) if you want to fit a constant background. "
        "Finally, 'suffix' must contain the names of each model, with the following "
        "format: [name1,name1,name3,...]. There must be as many suffixes as there are "
        "models in the 'models.fits' file.");
    print("Available options:");
    bullet("fitmask=...", "Must be a 2D FITS file containing the mask to define the fitting "
        "region. Only the pixels with a non-zero value in the mask will be used "
        "(default: all valid pixels are used).");
    bullet("residuals", "Set this flag if you want the program to generate residual cubes, "
        "removing the contribution of each component separately (default: do not "
        "generate residuals)");
    bullet("outdir", "Name of the directory into which the output files should be created. "
        "Default is the current directory.");
}
<|endoftext|>


<commit_before>#include "Parser.h"

#include <string>
#include <algorithm>    // copy

#include <boost/tokenizer.hpp>

void parseTSV(std::istream &in, std::function<void(Statement)> handler) {
    typedef boost::tokenizer< boost::char_separator<char> > Tokenizer;
    boost::char_separator<char> sep(" \t");

    std::vector< std::string > vec;
    std::string line;

    long id = 0;
    while (getline(in,line))
    {
        Tokenizer tok(line, sep);
        vec.assign(tok.begin(),tok.end());

        Statement stmt(id++,vec[0], vec[1], Value(vec[2]), qualifiers);
        statements[vec[0]];
        statements[vec[0]].push_back(stmt);
    }

}
<commit_msg>start draft of parser<commit_after>#include "Parser.h"

#include <regex>
#include <string>
#include <algorithm>    // copy

#include <boost/tokenizer.hpp>

PropertyValue parsePropertyValue(std::string property, value) {
    // regular expressions for the different value formats

    // QID of entities, properties, sources
    static std::regex re_entity("[QPS]\d*");

    // time format (e.g. +00000001967-01-17T00:00:00Z/11)
    static std::regex re_time("\\+\d{11}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}Z/\d{2}");

    // numbers (+/-1234.4567)
    static std::regex re_num("[+-]\d+(\\.\d+)?");

    // locations (@LAT/LONG)
    static std::regex re_loc("@\d{2}(\\.\d+)?/\d{3}(\\.\d+)?");

    if (std::regex_match(value, re_entity)) {

    }
}

void parseTSV(std::istream &in, std::function<void(Statement)> handler) {
    typedef boost::tokenizer< boost::escaped_list_error<char> > Tokenizer;
    boost::escaped_list_error<char> sep("\\"," \t","\"");

    std::vector< std::string > vec;
    std::string line;

    long id = 0;
    while (getline(in,line))
    {
        Tokenizer tok(line, sep);
        vec.assign(tok.begin(),tok.end());

        Statement::extensions_t qualifiers, sources;

        // 0, 1, 2 are subject, predicate and object; everything that follows
        // is a qualifier or source, depending on first letter of predicate.
        for(int i=3; i<vec.size(); i += 2) {
            PropertyValue p()

            if (vec[i][0] == 'S') {
                sources.push_back(p);
            } else {
                qualifiers.push_back(p);
            }
        }

        handler(Statement(-1, vec[0], PropertyValue(vec[1],vec[2]), qualifiers, sources, UNAPPROVED));
    }

}
<|endoftext|>


<commit_before>/*************************************************************************
 *                                                                       *
 * Open Dynamics Engine, Copyright (C) 2001-2003 Russell L. Smith.       *
 * All rights reserved.  Email: russ@q12.org   Web: www.q12.org          *
 *                                                                       *
 * This library is free software; you can redistribute it and/or         *
 * modify it under the terms of EITHER:                                  *
 *   (1) The GNU Lesser General Public License as published by the Free  *
 *       Software Foundation; either version 2.1 of the License, or (at  *
 *       your option) any later version. The text of the GNU Lesser      *
 *       General Public License is included with this library in the     *
 *       file LICENSE.TXT.                                               *
 *   (2) The BSD-style license that is included with this library in     *
 *       the file LICENSE-BSD.TXT.                                       *
 *                                                                       *
 * This library is distributed in the hope that it will be useful,       *
 * but WITHOUT ANY WARRANTY; without even the implied warranty of        *
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the files    *
 * LICENSE.TXT and LICENSE-BSD.TXT for more details.                     *
 *                                                                       *
 *************************************************************************/

// TriMesh code by Erwin de Vries.

#include <ode/collision.h>
#include <ode/matrix.h>
#include <ode/rotation.h>
#include <ode/odemath.h>

#if dTRIMESH_ENABLED

#include "collision_util.h"

#define TRIMESH_INTERNAL
#include "collision_trimesh_internal.h"

#if dTRIMESH_OPCODE
int dCollideRTL(dxGeom* g1, dxGeom* RayGeom, int Flags, dContactGeom* Contacts, int Stride){
	dxTriMesh* TriMesh = (dxTriMesh*)g1;

	const dVector3& TLPosition = *(const dVector3*)dGeomGetPosition(TriMesh);
	const dMatrix3& TLRotation = *(const dMatrix3*)dGeomGetRotation(TriMesh);

	RayCollider& Collider = TriMesh->_RayCollider;

	dReal Length = dGeomRayGetLength(RayGeom);

	int FirstContact, BackfaceCull;
	dGeomRayGetParams(RayGeom, &FirstContact, &BackfaceCull);
	int ClosestHit = dGeomRayGetClosestHit(RayGeom);

	Collider.SetFirstContact(FirstContact != 0);
	Collider.SetClosestHit(ClosestHit != 0);
	Collider.SetCulling(BackfaceCull != 0);
	Collider.SetMaxDist(Length);

	dVector3 Origin, Direction;
	dGeomRayGet(RayGeom, Origin, Direction);

	/* Make Ray */
	Ray WorldRay;
	WorldRay.mOrig.x = Origin[0];
	WorldRay.mOrig.y = Origin[1];
	WorldRay.mOrig.z = Origin[2];
	WorldRay.mDir.x = Direction[0];
	WorldRay.mDir.y = Direction[1];
	WorldRay.mDir.z = Direction[2];

	/* Intersect */
	Matrix4x4 amatrix;
        int TriCount = 0;
        if (Collider.Collide(WorldRay, TriMesh->Data->BVTree, &MakeMatrix(TLPosition, TLRotation, amatrix))) {
                TriCount = TriMesh->Faces.GetNbFaces();
        }

        if (TriCount == 0) {
                return 0;
        }

	const CollisionFace* Faces = TriMesh->Faces.GetFaces();

	int OutTriCount = 0;
	for (int i = 0; i < TriCount; i++) {
		if (OutTriCount == (Flags & 0xffff)) {
			break;
		}
		if (TriMesh->RayCallback == null ||
                    TriMesh->RayCallback(TriMesh, RayGeom, Faces[i].mFaceID,
                                         Faces[i].mU, Faces[i].mV)) {
			const int& TriIndex = Faces[i].mFaceID;
			if (!Callback(TriMesh, RayGeom, TriIndex)) {
                                continue;
                        }

			dContactGeom* Contact = SAFECONTACT(Flags, Contacts, OutTriCount, Stride);

			dVector3 dv[3];
			FetchTriangle(TriMesh, TriIndex, TLPosition, TLRotation, dv);

			float T = Faces[i].mDistance;
			Contact->pos[0] = Origin[0] + (Direction[0] * T);
			Contact->pos[1] = Origin[1] + (Direction[1] * T);
			Contact->pos[2] = Origin[2] + (Direction[2] * T);
			Contact->pos[3] = REAL(0.0);

			dVector3 vu;
			vu[0] = dv[1][0] - dv[0][0];
			vu[1] = dv[1][1] - dv[0][1];
			vu[2] = dv[1][2] - dv[0][2];
			vu[3] = REAL(0.0);

			dVector3 vv;
			vv[0] = dv[2][0] - dv[0][0];
			vv[1] = dv[2][1] - dv[0][1];
			vv[2] = dv[2][2] - dv[0][2];
			vv[3] = REAL(0.0);

			dCROSS(Contact->normal, =, vv, vu);	// Reversed

			dNormalize3(Contact->normal);

			Contact->depth = T;
			Contact->g1 = TriMesh;
			Contact->g2 = RayGeom;

			OutTriCount++;
		}
	}
	return OutTriCount;
}
#endif // dTRIMESH_OPCODE

#if dTRIMESH_GIMPACT
int dCollideRTL(dxGeom* g1, dxGeom* RayGeom, int Flags, dContactGeom* Contacts, int Stride)
{
	dxTriMesh* TriMesh = (dxTriMesh*)g1;

    dReal Length = dGeomRayGetLength(RayGeom);
	int FirstContact, BackfaceCull;
	dGeomRayGetParams(RayGeom, &FirstContact, &BackfaceCull);
	int ClosestHit = dGeomRayGetClosestHit(RayGeom);
	dVector3 Origin, Direction;
	dGeomRayGet(RayGeom, Origin, Direction);

    char intersect=0;
    GIM_TRIANGLE_RAY_CONTACT_DATA contact_data;

	if(ClosestHit)
	{
		intersect = gim_trimesh_ray_closest_collision(&TriMesh->m_collision_trimesh,Origin,Direction,Length,&contact_data);
	}
	else
	{
	    intersect = gim_trimesh_ray_collision(&TriMesh->m_collision_trimesh,Origin,Direction,Length,&contact_data);
	}

    if(intersect == 0)
	{
        return 0;
    }

	int OutTriCount = 0;

	if(TriMesh->RayCallback)
	{
        if(TriMesh->RayCallback(TriMesh, RayGeom, contact_data.m_face_id,
                                         contact_data.u , contact_data.v))
        {
            OutTriCount = 1;
        }
	}
	else
	{
	    OutTriCount = 1;
	}

	if(OutTriCount>0)
	{
	    dContactGeom* Contact = SAFECONTACT(Flags, Contacts, (OutTriCount-1), Stride);
        VEC_COPY(Contact->pos,contact_data.m_point);
        VEC_COPY(Contact->normal,contact_data.m_normal);
        Contact->depth = contact_data.tparam;
        Contact->g1 = TriMesh;
        Contact->g2 = RayGeom;
	}

	return OutTriCount;
}
#endif  // dTRIMESH_GIMPACT

#endif // dTRIMESH_ENABLED

<commit_msg>Improvement: expression optimized<commit_after>/*************************************************************************
 *                                                                       *
 * Open Dynamics Engine, Copyright (C) 2001-2003 Russell L. Smith.       *
 * All rights reserved.  Email: russ@q12.org   Web: www.q12.org          *
 *                                                                       *
 * This library is free software; you can redistribute it and/or         *
 * modify it under the terms of EITHER:                                  *
 *   (1) The GNU Lesser General Public License as published by the Free  *
 *       Software Foundation; either version 2.1 of the License, or (at  *
 *       your option) any later version. The text of the GNU Lesser      *
 *       General Public License is included with this library in the     *
 *       file LICENSE.TXT.                                               *
 *   (2) The BSD-style license that is included with this library in     *
 *       the file LICENSE-BSD.TXT.                                       *
 *                                                                       *
 * This library is distributed in the hope that it will be useful,       *
 * but WITHOUT ANY WARRANTY; without even the implied warranty of        *
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the files    *
 * LICENSE.TXT and LICENSE-BSD.TXT for more details.                     *
 *                                                                       *
 *************************************************************************/

// TriMesh code by Erwin de Vries.

#include <ode/collision.h>
#include <ode/matrix.h>
#include <ode/rotation.h>
#include <ode/odemath.h>

#if dTRIMESH_ENABLED

#include "collision_util.h"

#define TRIMESH_INTERNAL
#include "collision_trimesh_internal.h"

#if dTRIMESH_OPCODE
int dCollideRTL(dxGeom* g1, dxGeom* RayGeom, int Flags, dContactGeom* Contacts, int Stride){
	dxTriMesh* TriMesh = (dxTriMesh*)g1;

	const dVector3& TLPosition = *(const dVector3*)dGeomGetPosition(TriMesh);
	const dMatrix3& TLRotation = *(const dMatrix3*)dGeomGetRotation(TriMesh);

	RayCollider& Collider = TriMesh->_RayCollider;

	dReal Length = dGeomRayGetLength(RayGeom);

	int FirstContact, BackfaceCull;
	dGeomRayGetParams(RayGeom, &FirstContact, &BackfaceCull);
	int ClosestHit = dGeomRayGetClosestHit(RayGeom);

	Collider.SetFirstContact(FirstContact != 0);
	Collider.SetClosestHit(ClosestHit != 0);
	Collider.SetCulling(BackfaceCull != 0);
	Collider.SetMaxDist(Length);

	dVector3 Origin, Direction;
	dGeomRayGet(RayGeom, Origin, Direction);

	/* Make Ray */
	Ray WorldRay;
	WorldRay.mOrig.x = Origin[0];
	WorldRay.mOrig.y = Origin[1];
	WorldRay.mOrig.z = Origin[2];
	WorldRay.mDir.x = Direction[0];
	WorldRay.mDir.y = Direction[1];
	WorldRay.mDir.z = Direction[2];

	/* Intersect */
	Matrix4x4 amatrix;
        int TriCount = 0;
        if (Collider.Collide(WorldRay, TriMesh->Data->BVTree, &MakeMatrix(TLPosition, TLRotation, amatrix))) {
                TriCount = TriMesh->Faces.GetNbFaces();
        }

        if (TriCount == 0) {
                return 0;
        }

	const CollisionFace* Faces = TriMesh->Faces.GetFaces();

	int OutTriCount = 0;
	for (int i = 0; i < TriCount; i++) {
		if (OutTriCount == (Flags & 0xffff)) {
			break;
		}
		if (TriMesh->RayCallback == null ||
                    TriMesh->RayCallback(TriMesh, RayGeom, Faces[i].mFaceID,
                                         Faces[i].mU, Faces[i].mV)) {
			const int& TriIndex = Faces[i].mFaceID;
			if (!Callback(TriMesh, RayGeom, TriIndex)) {
                                continue;
                        }

			dContactGeom* Contact = SAFECONTACT(Flags, Contacts, OutTriCount, Stride);

			dVector3 dv[3];
			FetchTriangle(TriMesh, TriIndex, TLPosition, TLRotation, dv);

			float T = Faces[i].mDistance;
			Contact->pos[0] = Origin[0] + (Direction[0] * T);
			Contact->pos[1] = Origin[1] + (Direction[1] * T);
			Contact->pos[2] = Origin[2] + (Direction[2] * T);
			Contact->pos[3] = REAL(0.0);

			dVector3 vu;
			vu[0] = dv[1][0] - dv[0][0];
			vu[1] = dv[1][1] - dv[0][1];
			vu[2] = dv[1][2] - dv[0][2];
			vu[3] = REAL(0.0);

			dVector3 vv;
			vv[0] = dv[2][0] - dv[0][0];
			vv[1] = dv[2][1] - dv[0][1];
			vv[2] = dv[2][2] - dv[0][2];
			vv[3] = REAL(0.0);

			dCROSS(Contact->normal, =, vv, vu);	// Reversed

			dNormalize3(Contact->normal);

			Contact->depth = T;
			Contact->g1 = TriMesh;
			Contact->g2 = RayGeom;

			OutTriCount++;
		}
	}
	return OutTriCount;
}
#endif // dTRIMESH_OPCODE

#if dTRIMESH_GIMPACT
int dCollideRTL(dxGeom* g1, dxGeom* RayGeom, int Flags, dContactGeom* Contacts, int Stride)
{
	dxTriMesh* TriMesh = (dxTriMesh*)g1;

    dReal Length = dGeomRayGetLength(RayGeom);
	int FirstContact, BackfaceCull;
	dGeomRayGetParams(RayGeom, &FirstContact, &BackfaceCull);
	int ClosestHit = dGeomRayGetClosestHit(RayGeom);
	dVector3 Origin, Direction;
	dGeomRayGet(RayGeom, Origin, Direction);

    char intersect=0;
    GIM_TRIANGLE_RAY_CONTACT_DATA contact_data;

	if(ClosestHit)
	{
		intersect = gim_trimesh_ray_closest_collision(&TriMesh->m_collision_trimesh,Origin,Direction,Length,&contact_data);
	}
	else
	{
	    intersect = gim_trimesh_ray_collision(&TriMesh->m_collision_trimesh,Origin,Direction,Length,&contact_data);
	}

    if(intersect == 0)
	{
        return 0;
    }

	int OutTriCount = 0;

	if(!TriMesh->RayCallback ||
		TriMesh->RayCallback(TriMesh, RayGeom, contact_data.m_face_id, contact_data.u , contact_data.v))
	{
	    dContactGeom* Contact = SAFECONTACT(Flags, Contacts, (OutTriCount-1), Stride);
        VEC_COPY(Contact->pos,contact_data.m_point);
        VEC_COPY(Contact->normal,contact_data.m_normal);
        Contact->depth = contact_data.tparam;
        Contact->g1 = TriMesh;
        Contact->g2 = RayGeom;

		OutTriCount = 1;
	}

	return OutTriCount;
}
#endif  // dTRIMESH_GIMPACT

#endif // dTRIMESH_ENABLED

<|endoftext|>


<commit_before>////////////////////////////////////////////////////////////////////////////////
//
// Copyright 2006 - 2015, Paul Beckingham, Federico Hernandez.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included
// in all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
// OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
// THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
// SOFTWARE.
//
// http://www.opensource.org/licenses/mit-license.php
//
////////////////////////////////////////////////////////////////////////////////

#include <cmake.h>
#include <sstream>
#include <stdio.h>
#include <sys/ioctl.h>
#include <Context.h>
#include <main.h>
#include <text.h>
#include <i18n.h>

#ifdef SOLARIS
#include <sys/termios.h>
#endif

////////////////////////////////////////////////////////////////////////////////
int Context::getWidth ()
{
  // Determine window size.
  int width = config.getInteger ("defaultwidth");

  // A zero width value means 'infinity', which is approximated here by 2^16.
  if (width == 0)
    return 65536;

  if (config.getBoolean ("detection"))
  {
    if (terminal_width == 0 &&
        terminal_height == 0)
    {
      unsigned short buff[4];
      if (ioctl (fileno(stdout), TIOCGWINSZ, &buff) != -1)
      {
        terminal_height = buff[0];
        terminal_width = buff[1];
      }
    }

    width = terminal_width;

    // Ncurses does this, and perhaps we need to as well, to avoid a problem on
    // Cygwin where the display goes right up to the terminal width, and causes
    // and odd color wrapping problem.
    if (config.getBoolean ("avoidlastcolumn"))
      --width;
  }

  return width;
}

////////////////////////////////////////////////////////////////////////////////
int Context::getHeight ()
{
  // Determine window size.
  int height = config.getInteger ("defaultheight");

  // A zero height value means 'infinity', which is approximated here by 2^16.
  if (height == 0)
    return 65536;

  if (config.getBoolean ("detection"))
  {
    if (terminal_width == 0 &&
        terminal_height == 0)
    {
      unsigned short buff[4];
      if (ioctl (fileno(stdout), TIOCGWINSZ, &buff) != -1)
      {
        terminal_height = buff[0];
        terminal_width = buff[1];
      }
    }

    height = terminal_height;
  }

  return height;
}

////////////////////////////////////////////////////////////////////////////////
<commit_msg>Portability<commit_after>////////////////////////////////////////////////////////////////////////////////
//
// Copyright 2006 - 2015, Paul Beckingham, Federico Hernandez.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included
// in all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
// OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
// THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
// SOFTWARE.
//
// http://www.opensource.org/licenses/mit-license.php
//
////////////////////////////////////////////////////////////////////////////////

#include <cmake.h>
#include <sstream>
#include <stdio.h>
#include <unistd.h>
#include <sys/ioctl.h>
#include <Context.h>
#include <main.h>
#include <text.h>
#include <i18n.h>

#ifdef SOLARIS
#include <sys/termios.h>
#endif

////////////////////////////////////////////////////////////////////////////////
int Context::getWidth ()
{
  // Determine window size.
  int width = config.getInteger ("defaultwidth");

  // A zero width value means 'infinity', which is approximated here by 2^16.
  if (width == 0)
    return 65536;

  if (config.getBoolean ("detection"))
  {
    if (terminal_width == 0 &&
        terminal_height == 0)
    {
      unsigned short buff[4];
      if (ioctl (STDOUT_FILENO, TIOCGWINSZ, &buff) != -1)
      {
        terminal_height = buff[0];
        terminal_width = buff[1];
      }
    }

    width = terminal_width;

    // Ncurses does this, and perhaps we need to as well, to avoid a problem on
    // Cygwin where the display goes right up to the terminal width, and causes
    // and odd color wrapping problem.
    if (config.getBoolean ("avoidlastcolumn"))
      --width;
  }

  return width;
}

////////////////////////////////////////////////////////////////////////////////
int Context::getHeight ()
{
  // Determine window size.
  int height = config.getInteger ("defaultheight");

  // A zero height value means 'infinity', which is approximated here by 2^16.
  if (height == 0)
    return 65536;

  if (config.getBoolean ("detection"))
  {
    if (terminal_width == 0 &&
        terminal_height == 0)
    {
      unsigned short buff[4];
      if (ioctl (STDOUT_FILENO, TIOCGWINSZ, &buff) != -1)
      {
        terminal_height = buff[0];
        terminal_width = buff[1];
      }
    }

    height = terminal_height;
  }

  return height;
}

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


<commit_before>#include <iostream>
#include <string>
#include <llvm/Support/ErrorHandling.h>
#include "ast_printer.h"
#include "../ast/decl.h"
#include "../ast/stmt.h"

using namespace delta;

namespace {

int indentLevel = 0;

/// Inserts a line break followed by appropriate indentation.
std::ostream& br(std::ostream& out) {
    return out << '\n' << std::string(indentLevel * 4, ' ');
}

std::ostream& operator<<(std::ostream& out, const Expr& expr);
std::ostream& operator<<(std::ostream& out, const Stmt& stmt);

std::ostream& operator<<(std::ostream& out, const VariableExpr& expr) {
    return out << expr.identifier;
}

std::ostream& operator<<(std::ostream& out, const StrLiteralExpr& expr) {
    return out << '"' << expr.value << '"';
}

std::ostream& operator<<(std::ostream& out, const IntLiteralExpr& expr) {
    return out << expr.value;
}

std::ostream& operator<<(std::ostream& out, const BoolLiteralExpr& expr) {
    return out << (expr.value ? "true" : "false");
}

std::ostream& operator<<(std::ostream& out, const NullLiteralExpr& expr) {
    return out << "null";
}

std::ostream& operator<<(std::ostream& out, const ArrayLiteralExpr& expr) {
    out << "(array-literal";
    for (auto& e : expr.elements) out << " " << *e;
    return out << ")";
}

std::ostream& operator<<(std::ostream& out, const PrefixExpr& expr) {
    return out << "(" << expr.op << *expr.operand << ")";
}

std::ostream& operator<<(std::ostream& out, const BinaryExpr& expr) {
    return out << "(" << *expr.left << " " << expr.op << " " << *expr.right << ")";
}

std::ostream& operator<<(std::ostream& out, const CallExpr& expr) {
    out << "(call " << expr.funcName << " ";
    for (const Arg& arg : expr.args) {
        out << *arg.value;
        if (&arg != &expr.args.back()) out << " ";
    }
    return out << ")";
}

std::ostream& operator<<(std::ostream& out, const CastExpr& expr) {
    return out << "(cast " << *expr.expr << " " << expr.type << ")";
}

std::ostream& operator<<(std::ostream& out, const MemberExpr& expr) {
    return out << "(member-expr " << *expr.base << " " << expr.member << ")";
}

std::ostream& operator<<(std::ostream& out, const SubscriptExpr& expr) {
    return out << "(subscript " << *expr.array << " " << *expr.index << ")";
}

std::ostream& operator<<(std::ostream& out, const Expr& expr) {
    switch (expr.getKind()) {
        case ExprKind::VariableExpr:   return out << expr.getVariableExpr();
        case ExprKind::StrLiteralExpr: return out << expr.getStrLiteralExpr();
        case ExprKind::IntLiteralExpr: return out << expr.getIntLiteralExpr();
        case ExprKind::BoolLiteralExpr:return out << expr.getBoolLiteralExpr();
        case ExprKind::NullLiteralExpr:return out << expr.getNullLiteralExpr();
        case ExprKind::ArrayLiteralExpr:return out<< expr.getArrayLiteralExpr();
        case ExprKind::PrefixExpr:     return out << expr.getPrefixExpr();
        case ExprKind::BinaryExpr:     return out << expr.getBinaryExpr();
        case ExprKind::CallExpr:       return out << expr.getCallExpr();
        case ExprKind::CastExpr:       return out << expr.getCastExpr();
        case ExprKind::MemberExpr:     return out << expr.getMemberExpr();
        case ExprKind::SubscriptExpr:  return out << expr.getSubscriptExpr();
    }
}

std::ostream& operator<<(std::ostream& out, const ReturnStmt& stmt) {
    out << br << "(return-stmt ";
    for (const auto& value : stmt.values) {
        out << *value;
        if (&value != &stmt.values.back()) out << " ";
    }
    return out << ")";
}

std::ostream& operator<<(std::ostream& out, const VariableStmt& stmt) {
    return out << br << "(var-stmt " << stmt.decl->name << " " << *stmt.decl->initializer << ")";
}

std::ostream& operator<<(std::ostream& out, const IncrementStmt& stmt) {
    return out << br << "(inc-stmt " << *stmt.operand << ")";
}

std::ostream& operator<<(std::ostream& out, const DecrementStmt& stmt) {
    return out << br << "(dec-stmt " << *stmt.operand << ")";
}

std::ostream& operator<<(std::ostream& out, const DeferStmt& stmt) {
    return out << br << "(defer-stmt " << *stmt.expr << ")";
}

std::ostream& operator<<(std::ostream& out, const IfStmt& stmt) {
    out << br << "(if-stmt " << *stmt.condition;
    indentLevel++;
    out << br << "(then";
    indentLevel++;
    for (const auto& substmt : stmt.thenBody) {
        out << br << *substmt;
    }
    out << ")";
    indentLevel--;
    out << br << "(else";
    indentLevel++;
    for (const auto& substmt : stmt.elseBody) {
        if (!substmt->isIfStmt()) out << br;
        out << *substmt;
    }
    indentLevel--;
    out << ")";
    indentLevel--;
    return out << ")";
}

std::ostream& operator<<(std::ostream& out, const SwitchStmt& stmt) {
    out << br << "(switch-stmt " << *stmt.condition;
    indentLevel++;
    for (const SwitchCase& switchCase : stmt.cases) {
        out << br << "(case " << *switchCase.value;
        indentLevel++;
        for (const auto& substmt : switchCase.stmts) {
            out << br << *substmt;
        }
        indentLevel--;
        out << ")";
    }
    indentLevel--;
    return out << ")";
}

std::ostream& operator<<(std::ostream& out, const WhileStmt& stmt) {
    out << br << "(while-stmt " << *stmt.condition;
    indentLevel++;
    for (const auto& substmt : stmt.body) {
        out << br << *substmt;
    }
    indentLevel--;
    return out << ")";
}

std::ostream& operator<<(std::ostream& out, const BreakStmt&) {
    return out << "(break-stmt)";
}

std::ostream& operator<<(std::ostream& out, const AssignStmt& stmt) {
    return out << br << "(assign-stmt " << *stmt.lhs << " " << *stmt.rhs << ")";
}

std::ostream& operator<<(std::ostream& out, const AugAssignStmt& stmt) {
    return out << br << "(comp-assign-stmt " << stmt.op << " " << *stmt.lhs << " " << *stmt.rhs << ")";
}

std::ostream& operator<<(std::ostream& out, const Stmt& stmt) {
    switch (stmt.getKind()) {
        case StmtKind::ReturnStmt:    return out << stmt.getReturnStmt();
        case StmtKind::VariableStmt:  return out << stmt.getVariableStmt();
        case StmtKind::IncrementStmt: return out << stmt.getIncrementStmt();
        case StmtKind::DecrementStmt: return out << stmt.getDecrementStmt();
        case StmtKind::ExprStmt:      return out << *stmt.getExprStmt().expr;
        case StmtKind::DeferStmt:     return out << stmt.getDeferStmt();
        case StmtKind::IfStmt:        return out << stmt.getIfStmt();
        case StmtKind::SwitchStmt:    return out << stmt.getSwitchStmt();
        case StmtKind::WhileStmt:     return out << stmt.getWhileStmt();
        case StmtKind::BreakStmt:     return out << stmt.getBreakStmt();
        case StmtKind::AssignStmt:    return out << stmt.getAssignStmt();
        case StmtKind::AugAssignStmt: return out << stmt.getAugAssignStmt();
    }
}

std::ostream& operator<<(std::ostream& out, const ParamDecl& decl) {
    return out << "(" << decl.type << " " << decl.name << ")";
}

std::ostream& printFuncBase(std::ostream& out, const FuncDecl& decl) {
    out << " (";
    for (const ParamDecl& param : decl.params) {
        out << param;
        if (&param != &decl.params.back()) out << " ";
    }
    out << ") " << decl.returnType;

    if (!decl.isExtern()) {
        indentLevel++;
        for (const auto& stmt : *decl.body) {
            out << *stmt;
        }
        indentLevel--;
    }

    return out << ")";
}

std::ostream& operator<<(std::ostream& out, const FuncDecl& decl) {
    out << br << (decl.isExtern() ? "(extern-func-decl " : "(func-decl ");
    out << decl.name;
    return printFuncBase(out, decl);
}

std::ostream& operator<<(std::ostream& out, const GenericFuncDecl& decl) {
    out << br << (decl.func->isExtern() ? "(extern-generic-func-decl " : "(generic-func-decl ");
    out << decl.func->name << " (";
    for (const GenericParamDecl& genericParam : decl.genericParams) {
        out << genericParam.name;
        if (&genericParam != &decl.genericParams.back()) out << " ";
    }
    out << ")";
    return printFuncBase(out, *decl.func);
}

std::ostream& operator<<(std::ostream& out, const InitDecl& decl) {
    out << br << "(init-decl " << decl.getTypeDecl().name << " (";
    for (const ParamDecl& param : decl.params) {
        out << param;
        if (&param != &decl.params.back()) out << " ";
    }
    out << ")";
    indentLevel++;
    for (const auto& stmt : *decl.body) {
        out << *stmt;
    }
    indentLevel--;
    return out << ")";
}

std::ostream& operator<<(std::ostream& out, const DeinitDecl& decl) {
    out << br << "(deinit-decl " << decl.getTypeDecl().name;
    indentLevel++;
    for (const auto& stmt : *decl.body) {
        out << *stmt;
    }
    indentLevel--;
    return out << ")";
}

std::ostream& operator<<(std::ostream& out, const FieldDecl& decl) {
    return out << br << "(field-decl " << decl.type << " " << decl.name << ")";
}

std::ostream& operator<<(std::ostream& out, const TypeDecl& decl) {
    out << br << "(type-decl ";
    switch (decl.tag) {
        case TypeTag::Struct: out << "struct "; break;
        case TypeTag::Class: out << "class "; break;
    }
    out << decl.name;
    indentLevel++;
    for (const FieldDecl& field : decl.fields) {
        out << field;
    }
    indentLevel--;
    return out << ")";
}

std::ostream& operator<<(std::ostream& out, const VarDecl& decl) {
    return out << br << "(var-decl " << decl.name << " " << *decl.initializer << ")";
}

std::ostream& operator<<(std::ostream& out, const ImportDecl& decl) {
    return out << br << "(import-decl \"" << decl.target << "\")";
}

std::ostream& operator<<(std::ostream& out, const Decl& decl) {
    switch (decl.getKind()) {
        case DeclKind::ParamDecl: return out << decl.getParamDecl();
        case DeclKind::FuncDecl:  return out << decl.getFuncDecl();
        case DeclKind::GenericParamDecl: llvm_unreachable("handled via GenericFuncDecl");
        case DeclKind::GenericFuncDecl: return out << decl.getGenericFuncDecl();
        case DeclKind::InitDecl:  return out << decl.getInitDecl();
        case DeclKind::DeinitDecl:return out << decl.getDeinitDecl();
        case DeclKind::TypeDecl:  return out << decl.getTypeDecl();
        case DeclKind::VarDecl:   return out << decl.getVarDecl();
        case DeclKind::FieldDecl: return out << decl.getFieldDecl();
        case DeclKind::ImportDecl:return out << decl.getImportDecl();
    }
}

} // anonymous namespace

std::ostream& delta::operator<<(std::ostream& out,
                                const std::vector<std::unique_ptr<Decl>>& decls) {
    out << "(source-file";
    indentLevel++;
    for (const auto& decl : decls) {
        ::operator<<(out, *decl);
    }
    indentLevel--;
    return out << ")\n";
}
<commit_msg>Fix printing of "uninitialized" VariableStmt<commit_after>#include <iostream>
#include <string>
#include <llvm/Support/ErrorHandling.h>
#include "ast_printer.h"
#include "../ast/decl.h"
#include "../ast/stmt.h"

using namespace delta;

namespace {

int indentLevel = 0;

/// Inserts a line break followed by appropriate indentation.
std::ostream& br(std::ostream& out) {
    return out << '\n' << std::string(indentLevel * 4, ' ');
}

std::ostream& operator<<(std::ostream& out, const Expr& expr);
std::ostream& operator<<(std::ostream& out, const Stmt& stmt);

std::ostream& operator<<(std::ostream& out, const VariableExpr& expr) {
    return out << expr.identifier;
}

std::ostream& operator<<(std::ostream& out, const StrLiteralExpr& expr) {
    return out << '"' << expr.value << '"';
}

std::ostream& operator<<(std::ostream& out, const IntLiteralExpr& expr) {
    return out << expr.value;
}

std::ostream& operator<<(std::ostream& out, const BoolLiteralExpr& expr) {
    return out << (expr.value ? "true" : "false");
}

std::ostream& operator<<(std::ostream& out, const NullLiteralExpr& expr) {
    return out << "null";
}

std::ostream& operator<<(std::ostream& out, const ArrayLiteralExpr& expr) {
    out << "(array-literal";
    for (auto& e : expr.elements) out << " " << *e;
    return out << ")";
}

std::ostream& operator<<(std::ostream& out, const PrefixExpr& expr) {
    return out << "(" << expr.op << *expr.operand << ")";
}

std::ostream& operator<<(std::ostream& out, const BinaryExpr& expr) {
    return out << "(" << *expr.left << " " << expr.op << " " << *expr.right << ")";
}

std::ostream& operator<<(std::ostream& out, const CallExpr& expr) {
    out << "(call " << expr.funcName << " ";
    for (const Arg& arg : expr.args) {
        out << *arg.value;
        if (&arg != &expr.args.back()) out << " ";
    }
    return out << ")";
}

std::ostream& operator<<(std::ostream& out, const CastExpr& expr) {
    return out << "(cast " << *expr.expr << " " << expr.type << ")";
}

std::ostream& operator<<(std::ostream& out, const MemberExpr& expr) {
    return out << "(member-expr " << *expr.base << " " << expr.member << ")";
}

std::ostream& operator<<(std::ostream& out, const SubscriptExpr& expr) {
    return out << "(subscript " << *expr.array << " " << *expr.index << ")";
}

std::ostream& operator<<(std::ostream& out, const Expr& expr) {
    switch (expr.getKind()) {
        case ExprKind::VariableExpr:   return out << expr.getVariableExpr();
        case ExprKind::StrLiteralExpr: return out << expr.getStrLiteralExpr();
        case ExprKind::IntLiteralExpr: return out << expr.getIntLiteralExpr();
        case ExprKind::BoolLiteralExpr:return out << expr.getBoolLiteralExpr();
        case ExprKind::NullLiteralExpr:return out << expr.getNullLiteralExpr();
        case ExprKind::ArrayLiteralExpr:return out<< expr.getArrayLiteralExpr();
        case ExprKind::PrefixExpr:     return out << expr.getPrefixExpr();
        case ExprKind::BinaryExpr:     return out << expr.getBinaryExpr();
        case ExprKind::CallExpr:       return out << expr.getCallExpr();
        case ExprKind::CastExpr:       return out << expr.getCastExpr();
        case ExprKind::MemberExpr:     return out << expr.getMemberExpr();
        case ExprKind::SubscriptExpr:  return out << expr.getSubscriptExpr();
    }
}

std::ostream& operator<<(std::ostream& out, const ReturnStmt& stmt) {
    out << br << "(return-stmt ";
    for (const auto& value : stmt.values) {
        out << *value;
        if (&value != &stmt.values.back()) out << " ";
    }
    return out << ")";
}

std::ostream& operator<<(std::ostream& out, const VariableStmt& stmt) {
    out << br << "(var-stmt " << stmt.decl->name << " ";
    if (stmt.decl->initializer) out << *stmt.decl->initializer;
    else out << "uninitialized";
    return out << ")";
}

std::ostream& operator<<(std::ostream& out, const IncrementStmt& stmt) {
    return out << br << "(inc-stmt " << *stmt.operand << ")";
}

std::ostream& operator<<(std::ostream& out, const DecrementStmt& stmt) {
    return out << br << "(dec-stmt " << *stmt.operand << ")";
}

std::ostream& operator<<(std::ostream& out, const DeferStmt& stmt) {
    return out << br << "(defer-stmt " << *stmt.expr << ")";
}

std::ostream& operator<<(std::ostream& out, const IfStmt& stmt) {
    out << br << "(if-stmt " << *stmt.condition;
    indentLevel++;
    out << br << "(then";
    indentLevel++;
    for (const auto& substmt : stmt.thenBody) {
        out << br << *substmt;
    }
    out << ")";
    indentLevel--;
    out << br << "(else";
    indentLevel++;
    for (const auto& substmt : stmt.elseBody) {
        if (!substmt->isIfStmt()) out << br;
        out << *substmt;
    }
    indentLevel--;
    out << ")";
    indentLevel--;
    return out << ")";
}

std::ostream& operator<<(std::ostream& out, const SwitchStmt& stmt) {
    out << br << "(switch-stmt " << *stmt.condition;
    indentLevel++;
    for (const SwitchCase& switchCase : stmt.cases) {
        out << br << "(case " << *switchCase.value;
        indentLevel++;
        for (const auto& substmt : switchCase.stmts) {
            out << br << *substmt;
        }
        indentLevel--;
        out << ")";
    }
    indentLevel--;
    return out << ")";
}

std::ostream& operator<<(std::ostream& out, const WhileStmt& stmt) {
    out << br << "(while-stmt " << *stmt.condition;
    indentLevel++;
    for (const auto& substmt : stmt.body) {
        out << br << *substmt;
    }
    indentLevel--;
    return out << ")";
}

std::ostream& operator<<(std::ostream& out, const BreakStmt&) {
    return out << "(break-stmt)";
}

std::ostream& operator<<(std::ostream& out, const AssignStmt& stmt) {
    return out << br << "(assign-stmt " << *stmt.lhs << " " << *stmt.rhs << ")";
}

std::ostream& operator<<(std::ostream& out, const AugAssignStmt& stmt) {
    return out << br << "(comp-assign-stmt " << stmt.op << " " << *stmt.lhs << " " << *stmt.rhs << ")";
}

std::ostream& operator<<(std::ostream& out, const Stmt& stmt) {
    switch (stmt.getKind()) {
        case StmtKind::ReturnStmt:    return out << stmt.getReturnStmt();
        case StmtKind::VariableStmt:  return out << stmt.getVariableStmt();
        case StmtKind::IncrementStmt: return out << stmt.getIncrementStmt();
        case StmtKind::DecrementStmt: return out << stmt.getDecrementStmt();
        case StmtKind::ExprStmt:      return out << *stmt.getExprStmt().expr;
        case StmtKind::DeferStmt:     return out << stmt.getDeferStmt();
        case StmtKind::IfStmt:        return out << stmt.getIfStmt();
        case StmtKind::SwitchStmt:    return out << stmt.getSwitchStmt();
        case StmtKind::WhileStmt:     return out << stmt.getWhileStmt();
        case StmtKind::BreakStmt:     return out << stmt.getBreakStmt();
        case StmtKind::AssignStmt:    return out << stmt.getAssignStmt();
        case StmtKind::AugAssignStmt: return out << stmt.getAugAssignStmt();
    }
}

std::ostream& operator<<(std::ostream& out, const ParamDecl& decl) {
    return out << "(" << decl.type << " " << decl.name << ")";
}

std::ostream& printFuncBase(std::ostream& out, const FuncDecl& decl) {
    out << " (";
    for (const ParamDecl& param : decl.params) {
        out << param;
        if (&param != &decl.params.back()) out << " ";
    }
    out << ") " << decl.returnType;

    if (!decl.isExtern()) {
        indentLevel++;
        for (const auto& stmt : *decl.body) {
            out << *stmt;
        }
        indentLevel--;
    }

    return out << ")";
}

std::ostream& operator<<(std::ostream& out, const FuncDecl& decl) {
    out << br << (decl.isExtern() ? "(extern-func-decl " : "(func-decl ");
    out << decl.name;
    return printFuncBase(out, decl);
}

std::ostream& operator<<(std::ostream& out, const GenericFuncDecl& decl) {
    out << br << (decl.func->isExtern() ? "(extern-generic-func-decl " : "(generic-func-decl ");
    out << decl.func->name << " (";
    for (const GenericParamDecl& genericParam : decl.genericParams) {
        out << genericParam.name;
        if (&genericParam != &decl.genericParams.back()) out << " ";
    }
    out << ")";
    return printFuncBase(out, *decl.func);
}

std::ostream& operator<<(std::ostream& out, const InitDecl& decl) {
    out << br << "(init-decl " << decl.getTypeDecl().name << " (";
    for (const ParamDecl& param : decl.params) {
        out << param;
        if (&param != &decl.params.back()) out << " ";
    }
    out << ")";
    indentLevel++;
    for (const auto& stmt : *decl.body) {
        out << *stmt;
    }
    indentLevel--;
    return out << ")";
}

std::ostream& operator<<(std::ostream& out, const DeinitDecl& decl) {
    out << br << "(deinit-decl " << decl.getTypeDecl().name;
    indentLevel++;
    for (const auto& stmt : *decl.body) {
        out << *stmt;
    }
    indentLevel--;
    return out << ")";
}

std::ostream& operator<<(std::ostream& out, const FieldDecl& decl) {
    return out << br << "(field-decl " << decl.type << " " << decl.name << ")";
}

std::ostream& operator<<(std::ostream& out, const TypeDecl& decl) {
    out << br << "(type-decl ";
    switch (decl.tag) {
        case TypeTag::Struct: out << "struct "; break;
        case TypeTag::Class: out << "class "; break;
    }
    out << decl.name;
    indentLevel++;
    for (const FieldDecl& field : decl.fields) {
        out << field;
    }
    indentLevel--;
    return out << ")";
}

std::ostream& operator<<(std::ostream& out, const VarDecl& decl) {
    return out << br << "(var-decl " << decl.name << " " << *decl.initializer << ")";
}

std::ostream& operator<<(std::ostream& out, const ImportDecl& decl) {
    return out << br << "(import-decl \"" << decl.target << "\")";
}

std::ostream& operator<<(std::ostream& out, const Decl& decl) {
    switch (decl.getKind()) {
        case DeclKind::ParamDecl: return out << decl.getParamDecl();
        case DeclKind::FuncDecl:  return out << decl.getFuncDecl();
        case DeclKind::GenericParamDecl: llvm_unreachable("handled via GenericFuncDecl");
        case DeclKind::GenericFuncDecl: return out << decl.getGenericFuncDecl();
        case DeclKind::InitDecl:  return out << decl.getInitDecl();
        case DeclKind::DeinitDecl:return out << decl.getDeinitDecl();
        case DeclKind::TypeDecl:  return out << decl.getTypeDecl();
        case DeclKind::VarDecl:   return out << decl.getVarDecl();
        case DeclKind::FieldDecl: return out << decl.getFieldDecl();
        case DeclKind::ImportDecl:return out << decl.getImportDecl();
    }
}

} // anonymous namespace

std::ostream& delta::operator<<(std::ostream& out,
                                const std::vector<std::unique_ptr<Decl>>& decls) {
    out << "(source-file";
    indentLevel++;
    for (const auto& decl : decls) {
        ::operator<<(out, *decl);
    }
    indentLevel--;
    return out << ")\n";
}
<|endoftext|>


<commit_before>/* <x0/plugins/ssl/SslSocket.cpp>
 *
 * This file is part of the x0 web server project and is released under LGPL-3.
 * http://www.xzero.ws/
 *
 * (c) 2009-2010 Christian Parpart <trapni@gentoo.org>
 */

#include "SslSocket.h"
#include "SslDriver.h"
#include "SslContext.h"
#include <x0/Defines.h>
#include <x0/Buffer.h>
#include <x0/BufferRef.h>

#include <gnutls/gnutls.h>
#include <gnutls/extra.h>

#if 0
#	define TRACE(msg...) DEBUG("SslSocket: " msg)
#else
#	define TRACE(msg...)
#endif

SslSocket::SslSocket(SslDriver *driver, struct ev_loop *loop, int fd, int af) :
	x0::Socket(loop, fd, af),
#ifndef NDEBUG
	ctime_(ev_now(loop)),
#endif
	driver_(driver),
	context_(nullptr),
	session_()
{
	TRACE("SslSocket()");

	static int protocolPriorities_[] = { GNUTLS_TLS1_2, GNUTLS_TLS1_1, GNUTLS_TLS1_0, GNUTLS_SSL3, 0 };

	setSecure(true);
	setState(Handshake);

	gnutls_init(&session_, GNUTLS_SERVER);
	gnutls_protocol_set_priority(session_, protocolPriorities_);

	gnutls_handshake_set_post_client_hello_function(session_, &SslSocket::onClientHello);

	gnutls_certificate_server_set_request(session_, GNUTLS_CERT_REQUEST);
	gnutls_dh_set_prime_bits(session_, 1024);

	gnutls_session_enable_compatibility_mode(session_);

	gnutls_session_set_ptr(session_, this);
	gnutls_transport_set_ptr(session_, reinterpret_cast<gnutls_transport_ptr_t>(handle()));

	driver_->cache(this);
}

SslSocket::~SslSocket()
{
	TRACE("~SslSocket()");
	gnutls_deinit(session_);
}

int SslSocket::onClientHello(gnutls_session_t session)
{
	TRACE("onClientHello()");

	SslSocket *socket = (SslSocket *)gnutls_session_get_ptr(session);

	// find SNI server
	const int MAX_HOST_LEN = 255;
	std::size_t dataLen = MAX_HOST_LEN;
	char sniName[MAX_HOST_LEN];
	unsigned int sniType;

	int rv = gnutls_server_name_get(session, sniName, &dataLen, &sniType, 0);
	if (rv != 0)
	{
		TRACE("onClientHello(): gnutls_server_name_get() failed with (%d): %s", rv, gnutls_strerror(rv));

		// failed to get SNI from client, so try getting default context then.
		if (SslContext *cx = socket->driver_->selectContext("")) {
			cx->bind(socket);
			return 0;
		}

		return GNUTLS_E_UNIMPLEMENTED_FEATURE;
	}

	if (sniType != GNUTLS_NAME_DNS)
	{
		TRACE("onClientHello(): Unknown SNI type: %d", sniType);
		return GNUTLS_E_UNIMPLEMENTED_FEATURE;
	}

	TRACE("onClientHello(): SNI Name: \"%s\"", sniName);

	if (SslContext *cx = socket->driver_->selectContext(sniName))
		cx->bind(socket);

	return 0;
}

void SslSocket::handshake(int revents)
{
	TRACE("handshake(0x%04x)", revents);
	int rv = gnutls_handshake(session_);

	if (rv == GNUTLS_E_SUCCESS) {
		// handshake either completed or failed
		TRACE("SSL handshake complete. (time: %.4f)", ev_now(loop()) - ctime_);

		setState(Operational);
		setMode(Read);

		if (handshakeCallback_) {
			handshakeCallback_(this, handshakeData_);
		}
	} else if (rv != GNUTLS_E_AGAIN && rv != GNUTLS_E_INTERRUPTED) {
		TRACE("SSL handshake failed (%d): %s", rv, gnutls_strerror(rv));

		setState(Failure);

		if (handshakeCallback_) {
			handshakeCallback_(this, handshakeData_);
		}
	} else {
		TRACE("SSL partial handshake: (%d)", gnutls_record_get_direction(session_));
		switch (gnutls_record_get_direction(session_)) {
		case 0: // read
			setMode(Read);
			break;
		case 1: // write
			setMode(Write);
			break;
		default:
			break;
		}
	}
}

ssize_t SslSocket::read(x0::Buffer& result)
{
	if (result.size() == result.capacity())
		result.reserve(result.size() + 4096);

	ssize_t rv = gnutls_read(session_, result.end(), result.capacity() - result.size());
	if (rv < 0)
		return rv;

	result.resize(result.size() + rv);
	return rv;
}

ssize_t SslSocket::write(const void *buffer, size_t size)
{
	if (!size) {
		TRACE("SslSocket.write(empty buffer)");
		return 0;
	}

	ssize_t rv = gnutls_write(session_, buffer, size);
	TRACE("SslSocket.write(%ld bytes) = %ld", size, rv);
	if (rv >= 0)
		return rv;

	switch (rv) {
		case GNUTLS_E_AGAIN:
			errno = EAGAIN;
			break;
		case GNUTLS_E_INTERRUPTED:
			errno = EINTR;
			break;
		default:
			TRACE("gnutls_write error: %s", gnutls_strerror(rv));
			errno = EINVAL;
			setState(Failure);
			break;
	}

	return -1;
}

ssize_t SslSocket::write(int fd, off_t *offset, size_t nbytes)
{
	char buf[4096];
	ssize_t rv = pread(fd, buf, std::min(sizeof(buf), nbytes), *offset);

	if (rv > 0) {
		rv = write(buf, rv);

		if (rv > 0)
			*offset += rv;
	}

	return rv;
}
<commit_msg>[plugins] ssl: added some debug prints on gnutls function call errors<commit_after>/* <x0/plugins/ssl/SslSocket.cpp>
 *
 * This file is part of the x0 web server project and is released under LGPL-3.
 * http://www.xzero.ws/
 *
 * (c) 2009-2010 Christian Parpart <trapni@gentoo.org>
 */

#include "SslSocket.h"
#include "SslDriver.h"
#include "SslContext.h"
#include <x0/Defines.h>
#include <x0/Buffer.h>
#include <x0/BufferRef.h>

#include <gnutls/gnutls.h>
#include <gnutls/extra.h>

#if 0
#	define TRACE(msg...) DEBUG("SslSocket: " msg)
#else
#	define TRACE(msg...)
#endif

#define GNUTLS_CHECK(call) { \
	int rv = (call); \
	if (rv != GNUTLS_E_SUCCESS) { \
		TRACE("error running: %s = %d %s", #call, rv, gnutls_strerror(rv)); \
	} \
}

SslSocket::SslSocket(SslDriver *driver, struct ev_loop *loop, int fd, int af) :
	x0::Socket(loop, fd, af),
#ifndef NDEBUG
	ctime_(ev_now(loop)),
#endif
	driver_(driver),
	context_(nullptr),
	session_()
{
	TRACE("SslSocket()");

	static int protocolPriorities_[] = { GNUTLS_TLS1_2, GNUTLS_TLS1_1, GNUTLS_TLS1_0, GNUTLS_SSL3, 0 };

	setSecure(true);
	setState(Handshake);

	GNUTLS_CHECK( gnutls_init(&session_, GNUTLS_SERVER) );
	GNUTLS_CHECK( gnutls_protocol_set_priority(session_, protocolPriorities_) );

	gnutls_handshake_set_post_client_hello_function(session_, &SslSocket::onClientHello);

	gnutls_certificate_server_set_request(session_, GNUTLS_CERT_REQUEST);
	gnutls_dh_set_prime_bits(session_, 1024);

	gnutls_session_enable_compatibility_mode(session_);

	gnutls_session_set_ptr(session_, this);
	gnutls_transport_set_ptr(session_, reinterpret_cast<gnutls_transport_ptr_t>(handle()));

	driver_->cache(this);
}

SslSocket::~SslSocket()
{
	TRACE("~SslSocket()");
	gnutls_deinit(session_);
}

int SslSocket::onClientHello(gnutls_session_t session)
{
	TRACE("onClientHello()");

	SslSocket *socket = (SslSocket *)gnutls_session_get_ptr(session);

	// find SNI server
	const int MAX_HOST_LEN = 255;
	std::size_t dataLen = MAX_HOST_LEN;
	char sniName[MAX_HOST_LEN];
	unsigned int sniType;

	int rv = gnutls_server_name_get(session, sniName, &dataLen, &sniType, 0);
	if (rv != 0)
	{
		TRACE("onClientHello(): gnutls_server_name_get() failed with (%d): %s", rv, gnutls_strerror(rv));

		// failed to get SNI from client, so try getting default context then.
		if (SslContext *cx = socket->driver_->selectContext("")) {
			cx->bind(socket);
			return 0;
		}

		return GNUTLS_E_UNIMPLEMENTED_FEATURE;
	}

	if (sniType != GNUTLS_NAME_DNS)
	{
		TRACE("onClientHello(): Unknown SNI type: %d", sniType);
		return GNUTLS_E_UNIMPLEMENTED_FEATURE;
	}

	TRACE("onClientHello(): SNI Name: \"%s\"", sniName);

	if (SslContext *cx = socket->driver_->selectContext(sniName))
		cx->bind(socket);

	return 0;
}

void SslSocket::handshake(int revents)
{
	TRACE("handshake(0x%04x)", revents);
	int rv = gnutls_handshake(session_);

	if (rv == GNUTLS_E_SUCCESS) {
		// handshake either completed or failed
		TRACE("SSL handshake complete. (time: %.4f)", ev_now(loop()) - ctime_);

		setState(Operational);
		setMode(Read);

		if (handshakeCallback_) {
			handshakeCallback_(this, handshakeData_);
		}
	} else if (rv != GNUTLS_E_AGAIN && rv != GNUTLS_E_INTERRUPTED) {
		TRACE("SSL handshake failed (%d): %s", rv, gnutls_strerror(rv));

		setState(Failure);

		if (handshakeCallback_) {
			handshakeCallback_(this, handshakeData_);
		}
	} else {
		TRACE("SSL partial handshake: (%d)", gnutls_record_get_direction(session_));
		switch (gnutls_record_get_direction(session_)) {
		case 0: // read
			setMode(Read);
			break;
		case 1: // write
			setMode(Write);
			break;
		default:
			break;
		}
	}
}

ssize_t SslSocket::read(x0::Buffer& result)
{
	if (result.size() == result.capacity())
		result.reserve(result.size() + 4096);

	ssize_t rv = gnutls_read(session_, result.end(), result.capacity() - result.size());
	if (rv < 0)
		return rv;

	result.resize(result.size() + rv);
	return rv;
}

ssize_t SslSocket::write(const void *buffer, size_t size)
{
	if (!size) {
		TRACE("SslSocket.write(empty buffer)");
		return 0;
	}

	ssize_t rv = gnutls_write(session_, buffer, size);
	TRACE("SslSocket.write(%ld bytes) = %ld", size, rv);
	if (rv >= 0)
		return rv;

	switch (rv) {
		case GNUTLS_E_AGAIN:
			errno = EAGAIN;
			break;
		case GNUTLS_E_INTERRUPTED:
			errno = EINTR;
			break;
		default:
			TRACE("gnutls_write error: %s", gnutls_strerror(rv));
			errno = EINVAL;
			setState(Failure);
			break;
	}

	return -1;
}

ssize_t SslSocket::write(int fd, off_t *offset, size_t nbytes)
{
	char buf[4096];
	ssize_t rv = pread(fd, buf, std::min(sizeof(buf), nbytes), *offset);

	if (rv > 0) {
		rv = write(buf, rv);

		if (rv > 0)
			*offset += rv;
	}

	return rv;
}
<|endoftext|>