<commit_before>#include "WindowHelpers.h"

bool WindowHelpers::ToolTipCreated = false;
NOTIFYICONDATA WindowHelpers::Tray;

WindowHelpers::WindowHelpers() {
}

HWND WindowHelpers::getHandler() {
	return hwndWindow;
}

void WindowHelpers::TaskbarNotify(HWND hWnd) {
	if (!ToolTipCreated) {
		ShowWindow(hWnd, SHOW_WINDOW);

		if (IsWindowsVistaOrGreater()) {
			Tray.cbSize = sizeof(Tray);
			Tray.uVersion = NOTIFYICON_VERSION_4;
		}

		else if (IsWindowsXPOrGreater()) {
			Tray.cbSize = NOTIFYICONDATA_V3_SIZE;
			Tray.uVersion = NOTIFYICON_VERSION;
		}

		Shell_NotifyIcon(NIM_SETVERSION, &Tray);

		Tray.hIcon = LoadIcon(GetModuleHandle(NULL), MAKEINTRESOURCE(IDI_ICON1));
		Tray.hWnd = hWnd;
		strcpy_s(Tray.szTip, "NoCapsLock");
		Tray.uCallbackMessage = WM_CONTEXTMSGEVENT;
		Tray.uFlags = NIF_ICON | NIF_TIP | NIF_SHOWTIP | NIF_MESSAGE | NIF_GUID;
		Tray.uID = 01;
		Shell_NotifyIcon(NIM_ADD, &Tray);
		ToolTipCreated = true;
	}
}

int WindowHelpers::CreateWndProc() {
	HINSTANCE hInstance = (HINSTANCE)GetWindowLongPtr(helpers::GetConsoleWindow(), GWLP_HINSTANCE);

	WNDCLASS wc = { 0 };
	wc.hbrBackground = CreateSolidBrush(RGB(240, 240, 240));
	wc.hCursor = LoadCursor(hInstance, IDC_ARROW);
	wc.hIcon = LoadIcon(hInstance, IDI_APPLICATION);
	wc.hInstance = hInstance;
	wc.lpfnWndProc = WndProc;
	wc.lpszClassName = TEXT(NOCAPSCLASS);
	wc.style = CS_HREDRAW | CS_VREDRAW;

	if (!RegisterClass(&wc))
	{
		printf("Problem with WNDClass\n");
		return 1;
	}

	hwndWindow = CreateWindow(TEXT(NOCAPSCLASS),
		TEXT("Options - NoCapsLock"),
		WS_OVERLAPPED | WS_BORDER | WS_SYSMENU,
		520, 20, 550, 400,
		helpers::GetConsoleWindow(),
		NULL,
		hInstance, NULL);

	UpdateWindow(hwndWindow);
	TaskbarNotify(hwndWindow);

	MSG msg;
	while (GetMessage(&msg, hwndWindow, 0, 0))
	{
		TranslateMessage(&msg);
		DispatchMessage(&msg);
	}
	return 0;
}

void WindowHelpers::CreateUI(HWND hwnd) {
	std::map<DWORD_PTR, KeyObject::key_t> keys = KeyManager::GetKeyMap();
	std::map<DWORD_PTR, KeyObject::key_t>::iterator it;

	int startX = 2;
	for (it = keys.begin(); it != keys.end(); it++)
	{
		CreateWindow(TEXT("button"), TEXT(it->second.title.c_str()),
			WS_VISIBLE | WS_CHILD | BS_AUTOCHECKBOX,
			2, startX, 400, 20,
			hwnd, (HMENU) it->first, NULL, NULL);

		CheckDlgButton(hwnd, it->first, it->second.enabled ? BST_CHECKED : BST_UNCHECKED);
		startX += 22;
	}

	HFONT defaultFont;
	defaultFont = (HFONT)GetStockObject(DEFAULT_GUI_FONT);
	SendMessage(hwnd, WM_SETFONT, WPARAM(defaultFont), TRUE);
}

LRESULT CALLBACK WindowHelpers::WndProc(HWND hwnd, UINT message, WPARAM wparam, LPARAM lparam)
{
	switch (message) {
	case WM_CREATE:
		CreateUI(hwnd);
		break;
	case WM_CONTEXTMSGEVENT:
		switch (lparam) {
		case WM_RBUTTONUP:
			POINT cursor;
			GetCursorPos(&cursor);
			HMENU hMenu = LoadMenu(GetModuleHandle(NULL), MAKEINTRESOURCE(IDR_MENU1));

			std::map<DWORD_PTR, KeyObject::key_t> keys = KeyManager::GetKeyMap();
			std::map<DWORD_PTR, KeyObject::key_t>::iterator it;
			for (it = keys.begin(); it != keys.end(); it++)
			{
				InsertMenu(hMenu, 0, MFT_STRING, it->first, it->second.title.c_str());
				CheckMenuItem(hMenu, it->first, it->second.enabled ? MF_CHECKED : MF_UNCHECKED);
			}

#if _DEBUG
			InsertMenu(hMenu, 0, MFT_SEPARATOR, -1, "Show Options");
			InsertMenu(hMenu, 0, MFT_STRING, ID__SHOWOPTIONS, "Show Options");
#endif

			SetForegroundWindow(hwnd);
			TrackPopupMenu((HMENU)GetSubMenu(hMenu, 0), TPM_LEFTALIGN, cursor.x, cursor.y, 0, hwnd, NULL);
			break;
		}
		break;

	case WM_COMMAND:
		switch (wparam) {
		case ID__QUIT:
			PostMessage(helpers::GetConsoleWindow(), WM_CLOSE, 0, 0);
			PostQuitMessage(0);
			break;
		case ID__SOURCECODE:
			ShellExecute(0, 0, "https://github.com/Toyz/NoCapsLock", 0, 0, SW_SHOW);
			break;

		case ID__SHOWOPTIONS:
			ShowWindow(hwnd, 1);
			break;
		default:
			KeyObject::key_t KeyMeta = KeyManager::FindKey(wparam);

			if (KeyMeta.title != "") {
				KeyMeta.enabled = !KeyMeta.enabled;
				CheckDlgButton(hwnd, wparam, KeyMeta.enabled ? BST_CHECKED : BST_UNCHECKED);
				KeyManager::UpdateByKey(wparam, KeyMeta);
			}
		}

		break;
	case WM_CLOSE:
		//PostMessage(helpers::GetConsoleWindow(), WM_CLOSE, 0, 0);
		//PostQuitMessage(0);
		ShowWindow(hwnd, 0);
		return true;
	}

	return DefWindowProc(hwnd, message, wparam, lparam);
}

void WindowHelpers::ChangeTrayTitle(const char * title) {
	strcpy_s(Tray.szTip, title);
	Tray.uFlags = NIF_ICON | NIF_TIP | NIF_SHOWTIP | NIF_MESSAGE | NIF_GUID;
	Shell_NotifyIcon(NIM_MODIFY, &Tray);
}
<commit_msg>Fixed some casting warnings<commit_after>#include "WindowHelpers.h"

bool WindowHelpers::ToolTipCreated = false;
NOTIFYICONDATA WindowHelpers::Tray;

WindowHelpers::WindowHelpers() {
}

HWND WindowHelpers::getHandler() {
	return hwndWindow;
}

void WindowHelpers::TaskbarNotify(HWND hWnd) {
	if (!ToolTipCreated) {
		ShowWindow(hWnd, SHOW_WINDOW);

		if (IsWindowsVistaOrGreater()) {
			Tray.cbSize = sizeof(Tray);
			Tray.uVersion = NOTIFYICON_VERSION_4;
		}

		else if (IsWindowsXPOrGreater()) {
			Tray.cbSize = NOTIFYICONDATA_V3_SIZE;
			Tray.uVersion = NOTIFYICON_VERSION;
		}

		Shell_NotifyIcon(NIM_SETVERSION, &Tray);

		Tray.hIcon = LoadIcon(GetModuleHandle(NULL), MAKEINTRESOURCE(IDI_ICON1));
		Tray.hWnd = hWnd;
		strcpy_s(Tray.szTip, "NoCapsLock");
		Tray.uCallbackMessage = WM_CONTEXTMSGEVENT;
		Tray.uFlags = NIF_ICON | NIF_TIP | NIF_SHOWTIP | NIF_MESSAGE | NIF_GUID;
		Tray.uID = 01;
		Shell_NotifyIcon(NIM_ADD, &Tray);
		ToolTipCreated = true;
	}
}

int WindowHelpers::CreateWndProc() {
	HINSTANCE hInstance = (HINSTANCE)GetWindowLongPtr(helpers::GetConsoleWindow(), GWLP_HINSTANCE);

	WNDCLASS wc = { 0 };
	wc.hbrBackground = CreateSolidBrush(RGB(240, 240, 240));
	wc.hCursor = LoadCursor(hInstance, IDC_ARROW);
	wc.hIcon = LoadIcon(hInstance, IDI_APPLICATION);
	wc.hInstance = hInstance;
	wc.lpfnWndProc = WndProc;
	wc.lpszClassName = TEXT(NOCAPSCLASS);
	wc.style = CS_HREDRAW | CS_VREDRAW;

	if (!RegisterClass(&wc))
	{
		printf("Problem with WNDClass\n");
		return 1;
	}

	hwndWindow = CreateWindow(TEXT(NOCAPSCLASS),
		TEXT("Options - NoCapsLock"),
		WS_OVERLAPPED | WS_BORDER | WS_SYSMENU,
		520, 20, 550, 400,
		helpers::GetConsoleWindow(),
		NULL,
		hInstance, NULL);

	UpdateWindow(hwndWindow);
	TaskbarNotify(hwndWindow);

	MSG msg;
	while (GetMessage(&msg, hwndWindow, 0, 0))
	{
		TranslateMessage(&msg);
		DispatchMessage(&msg);
	}
	return 0;
}

void WindowHelpers::CreateUI(HWND hwnd) {
	std::map<DWORD_PTR, KeyObject::key_t> keys = KeyManager::GetKeyMap();
	std::map<DWORD_PTR, KeyObject::key_t>::iterator it;

	int startX = 2;
	for (it = keys.begin(); it != keys.end(); it++)
	{
		CreateWindow(TEXT("button"), TEXT(it->second.title.c_str()),
			WS_VISIBLE | WS_CHILD | BS_AUTOCHECKBOX,
			2, startX, 400, 20,
			hwnd, (HMENU) it->first, NULL, NULL);

		CheckDlgButton(hwnd, static_cast<int>(it->first), it->second.enabled ? BST_CHECKED : BST_UNCHECKED);
		startX += 22;
	}

	HFONT defaultFont;
	defaultFont = (HFONT)GetStockObject(DEFAULT_GUI_FONT);
	SendMessage(hwnd, WM_SETFONT, WPARAM(defaultFont), TRUE);
}

LRESULT CALLBACK WindowHelpers::WndProc(HWND hwnd, UINT message, WPARAM wparam, LPARAM lparam)
{
	switch (message) {
	case WM_CREATE:
		CreateUI(hwnd);
		break;
	case WM_CONTEXTMSGEVENT:
		switch (lparam) {
		case WM_RBUTTONUP:
			POINT cursor;
			GetCursorPos(&cursor);
			HMENU hMenu = LoadMenu(GetModuleHandle(NULL), MAKEINTRESOURCE(IDR_MENU1));

			std::map<DWORD_PTR, KeyObject::key_t> keys = KeyManager::GetKeyMap();
			std::map<DWORD_PTR, KeyObject::key_t>::iterator it;
			for (it = keys.begin(); it != keys.end(); it++)
			{
				InsertMenu(hMenu, 0, MFT_STRING, it->first, it->second.title.c_str());
				CheckMenuItem(hMenu, static_cast<int>(it->first), it->second.enabled ? MF_CHECKED : MF_UNCHECKED);
			}

#if _DEBUG
			InsertMenu(hMenu, 0, MFT_SEPARATOR, -1, "Show Options");
			InsertMenu(hMenu, 0, MFT_STRING, ID__SHOWOPTIONS, "Show Options");
#endif

			SetForegroundWindow(hwnd);
			TrackPopupMenu((HMENU)GetSubMenu(hMenu, 0), TPM_LEFTALIGN, cursor.x, cursor.y, 0, hwnd, NULL);
			break;
		}
		break;

	case WM_COMMAND:
		switch (wparam) {
		case ID__QUIT:
			PostMessage(helpers::GetConsoleWindow(), WM_CLOSE, 0, 0);
			PostQuitMessage(0);
			break;
		case ID__SOURCECODE:
			ShellExecute(0, 0, "https://github.com/Toyz/NoCapsLock", 0, 0, SW_SHOW);
			break;

		case ID__SHOWOPTIONS:
			ShowWindow(hwnd, 1);
			break;
		default:
			KeyObject::key_t KeyMeta = KeyManager::FindKey(wparam);

			if (KeyMeta.title != "") {
				KeyMeta.enabled = !KeyMeta.enabled;
				CheckDlgButton(hwnd, static_cast<int>(wparam), KeyMeta.enabled ? BST_CHECKED : BST_UNCHECKED);
				KeyManager::UpdateByKey(wparam, KeyMeta);
			}
		}

		break;
	case WM_CLOSE:
		//PostMessage(helpers::GetConsoleWindow(), WM_CLOSE, 0, 0);
		//PostQuitMessage(0);
		ShowWindow(hwnd, 0);
		return true;
	}

	return DefWindowProc(hwnd, message, wparam, lparam);
}

void WindowHelpers::ChangeTrayTitle(const char * title) {
	strcpy_s(Tray.szTip, title);
	Tray.uFlags = NIF_ICON | NIF_TIP | NIF_SHOWTIP | NIF_MESSAGE | NIF_GUID;
	Shell_NotifyIcon(NIM_MODIFY, &Tray);
}
<|endoftext|>


<commit_before>/**
 * @file
 * @brief Implementation of CorryvreckanWriter module
 * @copyright Copyright (c) 2017 CERN and the Allpix Squared authors.
 * This software is distributed under the terms of the MIT License, copied verbatim in the file "LICENSE.md".
 * In applying this license, CERN does not waive the privileges and immunities granted to it by virtue of its status as an
 * Intergovernmental Organization or submit itself to any jurisdiction.
 */

#include "CorryvreckanWriterModule.hpp"

#include <Math/RotationZYX.h>

#include <fstream>
#include <string>
#include <utility>

#include "core/utils/file.h"
#include "core/utils/log.h"

using namespace allpix;

CorryvreckanWriterModule::CorryvreckanWriterModule(Configuration& config, Messenger* messenger, GeometryManager* geoManager)
    : Module(config), messenger_(messenger), geometryManager_(geoManager) {

    // Require PixelCharge messages for single detector
    messenger_->bindMulti(this, &CorryvreckanWriterModule::pixel_messages_, MsgFlags::REQUIRED);

    config_.setDefault("file_name", "corryvreckanOutput.root");
    config_.setDefault("geometry_file", "corryvreckanGeometry.conf");
    config_.setDefault("output_mctruth", false);
}

// Set up the output trees
void CorryvreckanWriterModule::init() {

    // Check if MC data to be saved
    outputMCtruth_ = config_.get<bool>("output_mctruth", false);

    // Create output file and directories
    fileName_ = createOutputFile(allpix::add_file_extension(config_.get<std::string>("file_name"), "root"));
    outputFile_ = std::make_unique<TFile>(fileName_.c_str(), "RECREATE");
    outputFile_->cd();
    outputFile_->mkdir("pixels");
    if(outputMCtruth_) {
        outputFile_->mkdir("mcparticles");
    }

    // Create geometry file:
    geometryFileName_ = createOutputFile(allpix::add_file_extension(config_.get<std::string>("geometry_file"), "conf"));

    // Loop over all detectors and make trees for data
    auto detectors = geometryManager_->getDetectors();
    for(auto& detector : detectors) {

        // Get the detector ID and type
        std::string detectorID = detector->getName();

        // Create the tree
        std::string objectID = detectorID + "_pixels";
        std::string treeName = detectorID + "_Timepix3_pixels";
        outputTrees_[objectID] = new TTree(treeName.c_str(), treeName.c_str());
        outputTrees_[objectID]->Branch("time", &time_);

        // Map the pixel object to the tree
        treePixels_[objectID] = new corryvreckan::Pixel();
        outputTrees_[objectID]->Branch("pixels", &treePixels_[objectID]);

        // If MC truth needed then make trees for output
        if(!outputMCtruth_) {
            continue;
        }

        // Create the tree
        std::string objectID_MC = detectorID + "_mcparticles";
        std::string treeName_MC = detectorID + "_Timepix3_mcparticles";
        outputTreesMC_[objectID_MC] = new TTree(treeName_MC.c_str(), treeName_MC.c_str());
        outputTreesMC_[objectID_MC]->Branch("time", &time_);

        // Map the mc particle object to the tree
        treeMCParticles_[objectID_MC] = new corryvreckan::MCParticle();
        outputTreesMC_[objectID_MC]->Branch("mcparticles", &treeMCParticles_[objectID_MC]);
    }

    // Initialise the time
    time_ = 0;
}

// Make instantiations of Corryvreckan pixels, and store these in the trees during run time
void CorryvreckanWriterModule::run(unsigned int) {

    // Loop through all receieved messages
    for(auto& message : pixel_messages_) {

        auto detectorID = message->getDetector()->getName();
        auto objectID = detectorID + "_pixels";
        LOG(DEBUG) << "Receieved " << message->getData().size() << " pixel hits from detector " << detectorID;
        LOG(DEBUG) << "Time on event hits will be " << time_;

        // Loop through all pixels received
        for(auto& allpix_pixel : message->getData()) {

            // Make a new output pixel
            unsigned int pixelX = allpix_pixel.getPixel().getIndex().X();
            unsigned int pixelY = allpix_pixel.getPixel().getIndex().Y();
            double adc = allpix_pixel.getSignal();
            double time(time_);
            auto outputPixel = new corryvreckan::Pixel(detectorID, int(pixelY), int(pixelX), int(adc), time);

            LOG(DEBUG) << "Pixel (" << pixelX << "," << pixelY << ") written to device " << detectorID;

            // Map the pixel to the output tree and write it
            treePixels_[objectID] = outputPixel;
            outputTrees_[objectID]->Fill();
            delete outputPixel;

            // If writing MC truth then also write out associated particle info
            if(!outputMCtruth_) {
                continue;
            }

            // Get all associated particles
            auto mcp = allpix_pixel.getMCParticles();
            for(auto& particle : mcp) {

                // Create a new particle object
                std::string objectID_MC = detectorID + "_mcparticles";
                auto mcParticle = new corryvreckan::MCParticle(detectorID,
                                                               particle->getParticleID(),
                                                               particle->getLocalStartPoint(),
                                                               particle->getLocalEndPoint(),
                                                               time);

                // Map the mc particle to the output tree and write it
                treeMCParticles_[objectID_MC] = mcParticle;
                LOG(DEBUG) << "MC particle started locally at (" << mcParticle->getLocalStart().X() << ","
                           << mcParticle->getLocalStart().Y() << ") and ended at " << mcParticle->getLocalEnd().X() << ","
                           << mcParticle->getLocalEnd().Y() << ")";
                LOG(DEBUG) << "MC particle started globally at (" << particle->getGlobalStartPoint().X() << ","
                           << particle->getGlobalStartPoint().Y() << ") and ended at " << particle->getGlobalEndPoint().X()
                           << "," << particle->getGlobalEndPoint().Y() << ")";
                outputTreesMC_[objectID_MC]->Fill();
                delete mcParticle;
            }
        }
    }

    // Increment the time till the next event
    time_ += 10;
}

// Save the output trees to file
void CorryvreckanWriterModule::finalize() {

    // Loop over all detectors and store the trees
    auto detectors = geometryManager_->getDetectors();
    for(auto& detector : detectors) {

        // Get the detector ID and type
        std::string detectorID = detector->getName();
        std::string objectID = detectorID + "_pixels";

        // Move to the write output file
        outputFile_->cd();
        outputFile_->cd("pixels");
        outputTrees_[objectID]->Write();

        // Clean up the tree and remove object pointer
        delete outputTrees_[objectID];
        treePixels_[objectID] = nullptr;

        // Write the MC truth
        if(!outputMCtruth_) {
            continue;
        }

        std::string objectID_MC = detectorID + "_mcparticles";
        outputFile_->cd("mcparticles");
        outputTreesMC_[objectID_MC]->Write();

        // Clean up the tree and remove object pointer
        delete outputTreesMC_[objectID_MC];
        treeMCParticles_[objectID_MC] = nullptr;
    }
    outputFile_->Close();

    // Print statistics
    LOG(STATUS) << "Wrote output data to file:" << std::endl << fileName_;

    // Loop over all detectors and store the geometry:
    // Write geometry:
    std::ofstream geometry_file;
    if(!geometryFileName_.empty()) {
        geometry_file.open(geometryFileName_, std::ios_base::out | std::ios_base::trunc);
        if(!geometry_file.good()) {
            throw ModuleError("Cannot write to Corryvreckan geometry file");
        }

        geometry_file << "# Allpix Squared detector geometry - https://cern.ch/allpix-squared" << std::endl << std::endl;

        for(auto& detector : detectors) {
            geometry_file << "[" << detector->getName() << "]" << std::endl;
            geometry_file << "position = " << Units::display(detector->getPosition().x(), {"mm", "um"}) << ", "
                          << Units::display(detector->getPosition().y(), {"mm", "um"}) << ", "
                          << Units::display(detector->getPosition().z(), {"mm", "um"}) << std::endl;

            // Transform the rotation matrix to a ZYX rotation and invert it to get a XYZ rotation
            // This way we stay compatible to old Corryvreckan versions which only support XYZ.
            geometry_file << "orientation_mode = \"xyz\"" << std::endl;
            ROOT::Math::RotationZYX rotations(detector->getOrientation().Inverse());
            geometry_file << "orientation = " << Units::display(-rotations.Psi(), "deg") << ", "
                          << Units::display(-rotations.Theta(), "deg") << ", " << Units::display(-rotations.Phi(), "deg")
                          << std::endl;

            auto model = detector->getModel();
            geometry_file << "type = \"" << model->getType() << "\"" << std::endl;
            geometry_file << "pixel_pitch = " << Units::display(model->getPixelSize().x(), "um") << ", "
                          << Units::display(model->getPixelSize().y(), "um") << std::endl;
            geometry_file << "number_of_pixels = " << model->getNPixels().x() << ", " << model->getNPixels().y()
                          << std::endl;

            geometry_file << std::endl;
        }
    }
}
<commit_msg>CorryvreckanWriter: take actual detector model for the name, do not hardcode anything<commit_after>/**
 * @file
 * @brief Implementation of CorryvreckanWriter module
 * @copyright Copyright (c) 2017 CERN and the Allpix Squared authors.
 * This software is distributed under the terms of the MIT License, copied verbatim in the file "LICENSE.md".
 * In applying this license, CERN does not waive the privileges and immunities granted to it by virtue of its status as an
 * Intergovernmental Organization or submit itself to any jurisdiction.
 */

#include "CorryvreckanWriterModule.hpp"

#include <Math/RotationZYX.h>

#include <fstream>
#include <string>
#include <utility>

#include "core/utils/file.h"
#include "core/utils/log.h"

using namespace allpix;

CorryvreckanWriterModule::CorryvreckanWriterModule(Configuration& config, Messenger* messenger, GeometryManager* geoManager)
    : Module(config), messenger_(messenger), geometryManager_(geoManager) {

    // Require PixelCharge messages for single detector
    messenger_->bindMulti(this, &CorryvreckanWriterModule::pixel_messages_, MsgFlags::REQUIRED);

    config_.setDefault("file_name", "corryvreckanOutput.root");
    config_.setDefault("geometry_file", "corryvreckanGeometry.conf");
    config_.setDefault("output_mctruth", false);
}

// Set up the output trees
void CorryvreckanWriterModule::init() {

    // Check if MC data to be saved
    outputMCtruth_ = config_.get<bool>("output_mctruth", false);

    // Create output file and directories
    fileName_ = createOutputFile(allpix::add_file_extension(config_.get<std::string>("file_name"), "root"));
    outputFile_ = std::make_unique<TFile>(fileName_.c_str(), "RECREATE");
    outputFile_->cd();
    outputFile_->mkdir("pixels");
    if(outputMCtruth_) {
        outputFile_->mkdir("mcparticles");
    }

    // Create geometry file:
    geometryFileName_ = createOutputFile(allpix::add_file_extension(config_.get<std::string>("geometry_file"), "conf"));

    // Loop over all detectors and make trees for data
    auto detectors = geometryManager_->getDetectors();
    for(auto& detector : detectors) {

        // Get the detector ID and type
        std::string detectorID = detector->getName();
        std::string detectorModel = detector->getModel()->getType();

        // Create the tree
        std::string objectID = detectorID + "_pixels";
        std::string treeName = detectorID + "_" + detectorModel + "_pixels";
        outputTrees_[objectID] = new TTree(treeName.c_str(), treeName.c_str());
        outputTrees_[objectID]->Branch("time", &time_);

        // Map the pixel object to the tree
        treePixels_[objectID] = new corryvreckan::Pixel();
        outputTrees_[objectID]->Branch("pixels", &treePixels_[objectID]);

        // If MC truth needed then make trees for output
        if(!outputMCtruth_) {
            continue;
        }

        // Create the tree
        std::string objectID_MC = detectorID + "_mcparticles";
        std::string treeName_MC = detectorID + "_" + detectorModel + "_mcparticles";
        outputTreesMC_[objectID_MC] = new TTree(treeName_MC.c_str(), treeName_MC.c_str());
        outputTreesMC_[objectID_MC]->Branch("time", &time_);

        // Map the mc particle object to the tree
        treeMCParticles_[objectID_MC] = new corryvreckan::MCParticle();
        outputTreesMC_[objectID_MC]->Branch("mcparticles", &treeMCParticles_[objectID_MC]);
    }

    // Initialise the time
    time_ = 0;
}

// Make instantiations of Corryvreckan pixels, and store these in the trees during run time
void CorryvreckanWriterModule::run(unsigned int) {

    // Loop through all receieved messages
    for(auto& message : pixel_messages_) {

        auto detectorID = message->getDetector()->getName();
        auto objectID = detectorID + "_pixels";
        LOG(DEBUG) << "Receieved " << message->getData().size() << " pixel hits from detector " << detectorID;
        LOG(DEBUG) << "Time on event hits will be " << time_;

        // Loop through all pixels received
        for(auto& allpix_pixel : message->getData()) {

            // Make a new output pixel
            unsigned int pixelX = allpix_pixel.getPixel().getIndex().X();
            unsigned int pixelY = allpix_pixel.getPixel().getIndex().Y();
            double adc = allpix_pixel.getSignal();
            double time(time_);
            auto outputPixel = new corryvreckan::Pixel(detectorID, int(pixelY), int(pixelX), int(adc), time);

            LOG(DEBUG) << "Pixel (" << pixelX << "," << pixelY << ") written to device " << detectorID;

            // Map the pixel to the output tree and write it
            treePixels_[objectID] = outputPixel;
            outputTrees_[objectID]->Fill();
            delete outputPixel;

            // If writing MC truth then also write out associated particle info
            if(!outputMCtruth_) {
                continue;
            }

            // Get all associated particles
            auto mcp = allpix_pixel.getMCParticles();
            for(auto& particle : mcp) {

                // Create a new particle object
                std::string objectID_MC = detectorID + "_mcparticles";
                auto mcParticle = new corryvreckan::MCParticle(detectorID,
                                                               particle->getParticleID(),
                                                               particle->getLocalStartPoint(),
                                                               particle->getLocalEndPoint(),
                                                               time);

                // Map the mc particle to the output tree and write it
                treeMCParticles_[objectID_MC] = mcParticle;
                LOG(DEBUG) << "MC particle started locally at (" << mcParticle->getLocalStart().X() << ","
                           << mcParticle->getLocalStart().Y() << ") and ended at " << mcParticle->getLocalEnd().X() << ","
                           << mcParticle->getLocalEnd().Y() << ")";
                LOG(DEBUG) << "MC particle started globally at (" << particle->getGlobalStartPoint().X() << ","
                           << particle->getGlobalStartPoint().Y() << ") and ended at " << particle->getGlobalEndPoint().X()
                           << "," << particle->getGlobalEndPoint().Y() << ")";
                outputTreesMC_[objectID_MC]->Fill();
                delete mcParticle;
            }
        }
    }

    // Increment the time till the next event
    time_ += 10;
}

// Save the output trees to file
void CorryvreckanWriterModule::finalize() {

    // Loop over all detectors and store the trees
    auto detectors = geometryManager_->getDetectors();
    for(auto& detector : detectors) {

        // Get the detector ID and type
        std::string detectorID = detector->getName();
        std::string objectID = detectorID + "_pixels";

        // Move to the write output file
        outputFile_->cd();
        outputFile_->cd("pixels");
        outputTrees_[objectID]->Write();

        // Clean up the tree and remove object pointer
        delete outputTrees_[objectID];
        treePixels_[objectID] = nullptr;

        // Write the MC truth
        if(!outputMCtruth_) {
            continue;
        }

        std::string objectID_MC = detectorID + "_mcparticles";
        outputFile_->cd("mcparticles");
        outputTreesMC_[objectID_MC]->Write();

        // Clean up the tree and remove object pointer
        delete outputTreesMC_[objectID_MC];
        treeMCParticles_[objectID_MC] = nullptr;
    }
    outputFile_->Close();

    // Print statistics
    LOG(STATUS) << "Wrote output data to file:" << std::endl << fileName_;

    // Loop over all detectors and store the geometry:
    // Write geometry:
    std::ofstream geometry_file;
    if(!geometryFileName_.empty()) {
        geometry_file.open(geometryFileName_, std::ios_base::out | std::ios_base::trunc);
        if(!geometry_file.good()) {
            throw ModuleError("Cannot write to Corryvreckan geometry file");
        }

        geometry_file << "# Allpix Squared detector geometry - https://cern.ch/allpix-squared" << std::endl << std::endl;

        for(auto& detector : detectors) {
            geometry_file << "[" << detector->getName() << "]" << std::endl;
            geometry_file << "position = " << Units::display(detector->getPosition().x(), {"mm", "um"}) << ", "
                          << Units::display(detector->getPosition().y(), {"mm", "um"}) << ", "
                          << Units::display(detector->getPosition().z(), {"mm", "um"}) << std::endl;

            // Transform the rotation matrix to a ZYX rotation and invert it to get a XYZ rotation
            // This way we stay compatible to old Corryvreckan versions which only support XYZ.
            geometry_file << "orientation_mode = \"xyz\"" << std::endl;
            ROOT::Math::RotationZYX rotations(detector->getOrientation().Inverse());
            geometry_file << "orientation = " << Units::display(-rotations.Psi(), "deg") << ", "
                          << Units::display(-rotations.Theta(), "deg") << ", " << Units::display(-rotations.Phi(), "deg")
                          << std::endl;

            auto model = detector->getModel();
            geometry_file << "type = \"" << model->getType() << "\"" << std::endl;
            geometry_file << "pixel_pitch = " << Units::display(model->getPixelSize().x(), "um") << ", "
                          << Units::display(model->getPixelSize().y(), "um") << std::endl;
            geometry_file << "number_of_pixels = " << model->getNPixels().x() << ", " << model->getNPixels().y()
                          << std::endl;

            geometry_file << std::endl;
        }
    }
}
<|endoftext|>


<commit_before>// $Id$

AliEmcalClusterMaker* AddTaskEmcalClusterMaker(
  const UInt_t nonLinFunct   = AliEMCALRecoUtils::kBeamTestCorrected,
  const char *nClusters      = 0,
  const char *outClusName    = "EmcCaloClusters",
  const Double_t emin
  const Bool_t   histo       = kFALSE,
  const char *outputname     = "AnalysisResults.root"
)
{
  // Get the pointer to the existing analysis manager via the static access method.
  //==============================================================================
  AliAnalysisManager *mgr = AliAnalysisManager::GetAnalysisManager();
  if (!mgr)
  {
    ::Error("AddTaskHadCorr", "No analysis manager to connect to.");
    return NULL;
  }

  // Check the analysis type using the event handlers connected to the analysis manager.
  //==============================================================================
  if (!mgr->GetInputEventHandler())
  {
    ::Error("AddTaskHadCorr", "This task requires an input event handler");
    return NULL;
  }

  TString inputDataType = mgr->GetInputEventHandler()->GetDataType(); // can be "ESD" or "AOD"
  if (nClusters==0) {
    if (inputDataType != "ESD")
      nClusters = "caloClusters";
    else
      nClusters = "CaloClusters";
  }

  //-------------------------------------------------------
  // Init the task and do settings
  //-------------------------------------------------------

  TString name(Form("EmcalClusterMaker_%s", outClusName));
  AliEmcalClusterMaker *ecm = new AliEmcalClusterMaker(name, histo);
  ecm->SetOutClusName(outClusName);
  AliEMCALRecoUtils *ru = new AliEMCALRecoUtils;
  ru->SetNonLinearityFunction(nonLinFunct);
  ecm->SetRecoUtils(ru);
  AliClusterContainer *clusCont = ecm->AddClusterContainer(nClusters);
  //if (clusCont) clusCont->SetParticlePtCut(minPt);

  //-------------------------------------------------------
  // Final settings, pass to manager and set the containers
  //-------------------------------------------------------

  mgr->AddTask(ecm);

  // Create containers for input/output
  mgr->ConnectInput (ecm, 0, mgr->GetCommonInputContainer());

  if (histo) {
    AliAnalysisDataContainer *coecm = mgr->CreateContainer(name,
							   TList::Class(),
							   AliAnalysisManager::kOutputContainer,
							   outputname);
    mgr->ConnectOutput(ecm,1,coecm);
  }

  return ecm;
}
<commit_msg>use setter<commit_after>// $Id$

AliEmcalClusterMaker* AddTaskEmcalClusterMaker(
  const UInt_t nonLinFunct   = AliEMCALRecoUtils::kBeamTestCorrected,
  const char *nClusters      = 0,
  const char *outClusName    = "EmcCaloClusters",
  const Double_t emin        = 0.3,
  const Bool_t   histo       = kFALSE,
  const char *outputname     = "AnalysisResults.root"
)
{
  // Get the pointer to the existing analysis manager via the static access method.
  //==============================================================================
  AliAnalysisManager *mgr = AliAnalysisManager::GetAnalysisManager();
  if (!mgr)
  {
    ::Error("AddTaskHadCorr", "No analysis manager to connect to.");
    return NULL;
  }

  // Check the analysis type using the event handlers connected to the analysis manager.
  //==============================================================================
  if (!mgr->GetInputEventHandler())
  {
    ::Error("AddTaskHadCorr", "This task requires an input event handler");
    return NULL;
  }

  TString inputDataType = mgr->GetInputEventHandler()->GetDataType(); // can be "ESD" or "AOD"
  if (nClusters==0) {
    if (inputDataType != "ESD")
      nClusters = "caloClusters";
    else
      nClusters = "CaloClusters";
  }

  //-------------------------------------------------------
  // Init the task and do settings
  //-------------------------------------------------------

  TString name(Form("EmcalClusterMaker_%s", outClusName));
  AliEmcalClusterMaker *ecm = new AliEmcalClusterMaker(name, histo);
  ecm->SetOutClusName(outClusName);
  AliEMCALRecoUtils *ru = new AliEMCALRecoUtils;
  ru->SetNonLinearityFunction(nonLinFunct);
  ecm->SetRecoUtils(ru);
  AliClusterContainer *clusCont = ecm->AddClusterContainer(nClusters);
  clusCont->SetParticleECut(emin);

  //-------------------------------------------------------
  // Final settings, pass to manager and set the containers
  //-------------------------------------------------------

  mgr->AddTask(ecm);

  // Create containers for input/output
  mgr->ConnectInput (ecm, 0, mgr->GetCommonInputContainer());

  if (histo) {
    AliAnalysisDataContainer *coecm = mgr->CreateContainer(name,
							   TList::Class(),
							   AliAnalysisManager::kOutputContainer,
							   outputname);
    mgr->ConnectOutput(ecm,1,coecm);
  }

  return ecm;
}
<|endoftext|>


<commit_before>// $Id$
//
// Jet modelling task.
//
// Author: S.Aiola, C.Loizides

#include "AliJetModelBaseTask.h"

#include <TClonesArray.h>
#include <TF1.h>
#include <TLorentzVector.h>
#include <TRandom3.h>

#include "AliAODCaloCluster.h"
#include "AliAnalysisManager.h"
#include "AliEMCALDigit.h"
#include "AliEMCALGeometry.h"
#include "AliEMCALRecPoint.h"
#include "AliESDCaloCluster.h"
#include "AliLog.h"
#include "AliPicoTrack.h"
#include "AliVCluster.h"
#include "AliVEvent.h"

ClassImp(AliJetModelBaseTask)

//________________________________________________________________________
AliJetModelBaseTask::AliJetModelBaseTask() :
  AliAnalysisTaskSE("AliJetModelBaseTask"),
  fGeomName(),
  fTracksName(),
  fOutTracksName(),
  fCaloName(),
  fOutCaloName(),
  fSuffix(),
  fEtaMin(-1),
  fEtaMax(1),
  fPhiMin(0),
  fPhiMax(TMath::Pi() * 2),
  fPtMin(0),
  fPtMax(0),
  fCopyArray(kTRUE),
  fNClusters(0),
  fNTracks(0),
  fMarkMC(kTRUE),
  fPtSpectrum(0),
  fIsInit(0),
  fGeom(0),
  fClusters(0),
  fOutClusters(0),
  fTracks(0),
  fOutTracks(0)
{
  // Default constructor.
}

//________________________________________________________________________
AliJetModelBaseTask::AliJetModelBaseTask(const char *name) :
  AliAnalysisTaskSE(name),
  fGeomName(""),
  fTracksName("PicoTracks"),
  fOutTracksName("PicoTracksEmbedded"),
  fCaloName("CaloClustersCorr"),
  fOutCaloName("CaloClustersCorrEmbedded"),
  fSuffix("Processed"),
  fEtaMin(-1),
  fEtaMax(1),
  fPhiMin(0),
  fPhiMax(TMath::Pi() * 2),
  fPtMin(50),
  fPtMax(60),
  fCopyArray(kTRUE),
  fNClusters(0),
  fNTracks(1),
  fMarkMC(kTRUE),
  fPtSpectrum(0),
  fIsInit(0),
  fGeom(0),
  fClusters(0),
  fOutClusters(0),
  fTracks(0),
  fOutTracks(0)
{
  // Standard constructor.
}

//________________________________________________________________________
AliJetModelBaseTask::~AliJetModelBaseTask()
{
  // Destructor
}

//________________________________________________________________________
void AliJetModelBaseTask::UserExec(Option_t *)
{
  // Execute per event.

  if (!fIsInit)
    ExecOnce();
  Run();
}

//________________________________________________________________________
AliVCluster* AliJetModelBaseTask::AddCluster(Double_t e, Double_t eta, Double_t phi)
{
  // Add a cluster to the event.

  Int_t absId = 0;
  if (eta < -100 || phi < 0) {
    GetRandomCell(eta, phi, absId);
  }
  else {
    fGeom->EtaPhiFromIndex(absId, eta, phi);
  }

  if (absId == -1) {
    AliWarning(Form("Unable to embed cluster in eta = %f, phi = %f!"
		    " Maybe the eta-phi range is not inside the EMCal acceptance (eta = [%f, %f], phi = [%f, %f])",
		    eta, phi, fEtaMin, fEtaMax, fPhiMin, fPhiMax));
    return 0;
  }

  if (e < 0) {
    Double_t pt = GetRandomPt();
    TLorentzVector nPart;
    nPart.SetPtEtaPhiM(pt, eta, phi, 0);
    e = nPart.E();
  }

  return AddCluster(e, absId);
}

//________________________________________________________________________
AliVCluster* AliJetModelBaseTask::AddCluster(Double_t e, Int_t absId)
{
  // Add a cluster to the event.

  const Int_t nClusters = fOutClusters->GetEntriesFast();

  TClonesArray digits("AliEMCALDigit", 1);

  AliEMCALDigit *digit = static_cast<AliEMCALDigit*>(digits.New(0));
  digit->SetId(absId);
  digit->SetIndexInList(0);
  digit->SetType(AliEMCALDigit::kHG);
  digit->SetAmplitude(e);

  AliEMCALRecPoint recPoint("");
  recPoint.AddDigit(*digit, e, kFALSE);
  recPoint.EvalGlobalPosition(0, &digits);

  TVector3 gpos;
  recPoint.GetGlobalPosition(gpos);
  Float_t g[3];
  gpos.GetXYZ(g);

  AliVCluster *cluster = static_cast<AliVCluster*>(fOutClusters->New(nClusters));
  cluster->SetType(AliVCluster::kEMCALClusterv1);
  cluster->SetE(recPoint.GetEnergy());
  cluster->SetPosition(g);
  cluster->SetNCells(1);
  UShort_t shortAbsId = absId;
  cluster->SetCellsAbsId(&shortAbsId);
  Double32_t fract = 1;
  cluster->SetCellsAmplitudeFraction(&fract);
  cluster->SetID(nClusters);
  cluster->SetEmcCpvDistance(-1);
  if (fMarkMC)
    cluster->SetChi2(100); // MC flag!

  return cluster;
}

//________________________________________________________________________
AliPicoTrack* AliJetModelBaseTask::AddTrack(Double_t pt, Double_t eta, Double_t phi)
{
  // Add a track to the event.

  const Int_t nTracks = fOutTracks->GetEntriesFast();

  if (pt < 0)
    pt = GetRandomPt();
  if (eta < -100)
    eta = GetRandomEta();
  if (phi < 0)
    phi = GetRandomPhi();

  Int_t label = fMarkMC ? 100 : 0;

  AliPicoTrack *track = new ((*fOutTracks)[nTracks]) AliPicoTrack(pt,
								  eta,
								  phi,
								  1,
								  label,    // MC flag!
								  0,
								  0,
								  kFALSE);
  return track;
}

//________________________________________________________________________
void AliJetModelBaseTask::CopyClusters()
{
  // Copy all the clusters in the new collection

  Bool_t esdMode = (Bool_t)(fClusters->GetClass()->GetBaseClass("AliESDCaloCluster") != 0);
  const Int_t nClusters = fClusters->GetEntriesFast();

  if (esdMode) {
    for (Int_t i = 0; i < nClusters; ++i) {
      AliESDCaloCluster *esdcluster = static_cast<AliESDCaloCluster*>(fClusters->At(i));
      if (!esdcluster)
	continue;
      if (!esdcluster->IsEMCAL())
	continue;
      new ((*fOutClusters)[i]) AliESDCaloCluster(*esdcluster);
    }
  }
  else {
    for (Int_t i = 0; i < nClusters; ++i) {
      AliAODCaloCluster *aodcluster = static_cast<AliAODCaloCluster*>(fClusters->At(i));
      if (!aodcluster)
	continue;
      if (!aodcluster->IsEMCAL())
	continue;
      new ((*fOutClusters)[i]) AliAODCaloCluster(*aodcluster);
    }
  }
}

//________________________________________________________________________
void AliJetModelBaseTask::CopyTracks()
{
  const Int_t nTracks = fTracks->GetEntriesFast();
  for (Int_t i = 0; i < nTracks; ++i) {
    AliPicoTrack *track = static_cast<AliPicoTrack*>(fTracks->At(i));
    if (!track)
      continue;
    new ((*fOutTracks)[i]) AliPicoTrack(*track);
  }
}

//________________________________________________________________________
void AliJetModelBaseTask::ExecOnce()
{
  // Init task.

  if (fPtMax < fPtMin) {
    AliWarning (Form("PtMax (%f) < PtMin (%f), setting PtMax = PtMin = %f", fPtMax, fPtMin, fPtMin));
    fPtMax = fPtMin;
  }

  if (fEtaMax < fEtaMin) {
    AliWarning (Form("EtaMax (%f) < EtaMin (%f), setting EtaMax = EtaMin = %f", fEtaMax, fEtaMin, fEtaMin));
    fEtaMax = fEtaMin;
  }

  if (fPhiMax < fPhiMin) {
    AliWarning (Form("PhiMax (%f) < PhiMin (%f), setting PhiMax = PhiMin = %f", fPhiMax, fPhiMin, fPhiMin));
    fPhiMax = fPhiMin;
  }

  if (fNTracks > 0 && !fTracks && !fTracksName.IsNull()) {
    fTracks = dynamic_cast<TClonesArray*>(InputEvent()->FindListObject(fTracksName));
    if (!fTracks) {
      AliError(Form("%s: Couldn't retrieve tracks with name %s!", GetName(), fTracksName.Data()));
      return;
    }

    if (!fTracks->GetClass()->GetBaseClass("AliPicoTrack")) {
      AliError(Form("%s: Collection %s does not contain AliPicoTrack objects!", GetName(), fTracksName.Data()));
      fTracks = 0;
      return;
    }

    if (!fOutTracks) {
      fOutTracksName = fTracksName;
      if (fCopyArray) {
	fOutTracksName += fSuffix;
	fOutTracks = new TClonesArray("AliPicoTrack", fTracks->GetSize());
	fOutTracks->SetName(fOutTracksName);
      }
      else {
	fOutTracks = fTracks;
      }
    }

    if (fCopyArray) {
      if (!(InputEvent()->FindListObject(fOutTracksName)))
	InputEvent()->AddObject(fOutTracks);
    }
  }

  if (fNClusters > 0 && !fClusters && !fCaloName.IsNull()) {
    fClusters = dynamic_cast<TClonesArray*>(InputEvent()->FindListObject(fCaloName));

    if (!fClusters) {
      AliError(Form("%s: Couldn't retrieve clusters with name %s!", GetName(), fCaloName.Data()));
      return;
    }

    if (!fClusters->GetClass()->GetBaseClass("AliVCluster")) {
      AliError(Form("%s: Collection %s does not contain AliVCluster objects!", GetName(), fCaloName.Data()));
      fClusters = 0;
      return;
    }

    if (!fOutClusters) {
      fOutCaloName = fCaloName;
      if (fCopyArray) {
	fOutCaloName += fSuffix;
	fOutClusters = new TClonesArray(fClusters->GetClass()->GetName(), fClusters->GetSize());
	fOutClusters->SetName(fOutCaloName);
      }
      else {
	fOutClusters = fClusters;
      }
    }

    if (fCopyArray) {
      if (!(InputEvent()->FindListObject(fOutCaloName)))
	InputEvent()->AddObject(fOutClusters);
    }

    if (!fGeom) {
      if (fGeomName.Length() > 0) {
	fGeom = AliEMCALGeometry::GetInstance(fGeomName);
	if (!fGeom)
	  AliError(Form("Could not get geometry with name %s!", fGeomName.Data()));
      } else {
	fGeom = AliEMCALGeometry::GetInstance();
	if (!fGeom)
	  AliError("Could not get default geometry!");
      }
    }

    const Double_t EmcalMinEta = fGeom->GetArm1EtaMin();
    const Double_t EmcalMaxEta = fGeom->GetArm1EtaMax();
    const Double_t EmcalMinPhi = fGeom->GetArm1PhiMin() * TMath::DegToRad();
    const Double_t EmcalMaxPhi = fGeom->GetArm1PhiMax() * TMath::DegToRad();

    if (fEtaMax > EmcalMaxEta) fEtaMax = EmcalMaxEta;
    if (fEtaMax < EmcalMinEta) fEtaMax = EmcalMinEta;
    if (fEtaMin > EmcalMaxEta) fEtaMin = EmcalMaxEta;
    if (fEtaMin < EmcalMinEta) fEtaMin = EmcalMinEta;

    if (fPhiMax > EmcalMaxPhi) fPhiMax = EmcalMaxPhi;
    if (fPhiMax < EmcalMinPhi) fPhiMax = EmcalMinPhi;
    if (fPhiMin > EmcalMaxPhi) fPhiMin = EmcalMaxPhi;
    if (fPhiMin < EmcalMinPhi) fPhiMin = EmcalMinPhi;
  }
}

//________________________________________________________________________
void AliJetModelBaseTask::GetRandomCell(Double_t &eta, Double_t &phi, Int_t &absId)
{
  // Get random cell.

  Int_t repeats = 0;
  Double_t rndEta = eta;
  Double_t rndPhi = phi;
  do {
    if (eta < -100)
      rndEta = GetRandomEta();
    if (phi < 0)
      rndPhi = GetRandomPhi();
    fGeom->GetAbsCellIdFromEtaPhi(rndEta, rndPhi, absId);
    repeats++;
  } while (absId == -1 && repeats < 100);

  if (!(absId > -1)) {
    AliWarning(Form("Could not extract random cluster! Random eta-phi extracted more than 100 times!\n"
		    "eta [%f, %f], phi [%f, %f]\n", fEtaMin, fEtaMax, fPhiMin, fPhiMax));
  }
  else {
    eta = rndEta;
    phi = rndPhi;
  }
}

//________________________________________________________________________
Double_t AliJetModelBaseTask::GetRandomEta()
{
  // Get random eta.

  return gRandom->Rndm() * (fEtaMax - fEtaMin) + fEtaMin;
}

//________________________________________________________________________
Double_t AliJetModelBaseTask::GetRandomPhi()
{
  // Get random phi.

  return gRandom->Rndm() * (fPhiMax - fPhiMin) + fPhiMin;
}

//________________________________________________________________________
Double_t AliJetModelBaseTask::GetRandomPt()
{
  // Get random pt.

  if (fPtSpectrum)
    return fPtSpectrum->GetRandom();
  else
    return gRandom->Rndm() * (fPtMax - fPtMin) + fPtMin;
}

//________________________________________________________________________
void AliJetModelBaseTask::Run()
{
  // Run.
}
<commit_msg>fix init once<commit_after>// $Id$
//
// Jet modelling task.
//
// Author: S.Aiola, C.Loizides

#include "AliJetModelBaseTask.h"

#include <TClonesArray.h>
#include <TF1.h>
#include <TLorentzVector.h>
#include <TRandom3.h>

#include "AliAODCaloCluster.h"
#include "AliAnalysisManager.h"
#include "AliEMCALDigit.h"
#include "AliEMCALGeometry.h"
#include "AliEMCALRecPoint.h"
#include "AliESDCaloCluster.h"
#include "AliLog.h"
#include "AliPicoTrack.h"
#include "AliVCluster.h"
#include "AliVEvent.h"

ClassImp(AliJetModelBaseTask)

//________________________________________________________________________
AliJetModelBaseTask::AliJetModelBaseTask() :
  AliAnalysisTaskSE("AliJetModelBaseTask"),
  fGeomName(),
  fTracksName(),
  fOutTracksName(),
  fCaloName(),
  fOutCaloName(),
  fSuffix(),
  fEtaMin(-1),
  fEtaMax(1),
  fPhiMin(0),
  fPhiMax(TMath::Pi() * 2),
  fPtMin(0),
  fPtMax(0),
  fCopyArray(kTRUE),
  fNClusters(0),
  fNTracks(0),
  fMarkMC(kTRUE),
  fPtSpectrum(0),
  fIsInit(0),
  fGeom(0),
  fClusters(0),
  fOutClusters(0),
  fTracks(0),
  fOutTracks(0)
{
  // Default constructor.
}

//________________________________________________________________________
AliJetModelBaseTask::AliJetModelBaseTask(const char *name) :
  AliAnalysisTaskSE(name),
  fGeomName(""),
  fTracksName("PicoTracks"),
  fOutTracksName("PicoTracksEmbedded"),
  fCaloName("CaloClustersCorr"),
  fOutCaloName("CaloClustersCorrEmbedded"),
  fSuffix("Processed"),
  fEtaMin(-1),
  fEtaMax(1),
  fPhiMin(0),
  fPhiMax(TMath::Pi() * 2),
  fPtMin(50),
  fPtMax(60),
  fCopyArray(kTRUE),
  fNClusters(0),
  fNTracks(1),
  fMarkMC(kTRUE),
  fPtSpectrum(0),
  fIsInit(0),
  fGeom(0),
  fClusters(0),
  fOutClusters(0),
  fTracks(0),
  fOutTracks(0)
{
  // Standard constructor.
}

//________________________________________________________________________
AliJetModelBaseTask::~AliJetModelBaseTask()
{
  // Destructor
}

//________________________________________________________________________
void AliJetModelBaseTask::UserExec(Option_t *)
{
  // Execute per event.

  if (!fIsInit) {
    ExecOnce();
    fIsInit = 1;
  }
  Run();
}

//________________________________________________________________________
AliVCluster* AliJetModelBaseTask::AddCluster(Double_t e, Double_t eta, Double_t phi)
{
  // Add a cluster to the event.

  Int_t absId = 0;
  if (eta < -100 || phi < 0) {
    GetRandomCell(eta, phi, absId);
  }
  else {
    fGeom->EtaPhiFromIndex(absId, eta, phi);
  }

  if (absId == -1) {
    AliWarning(Form("Unable to embed cluster in eta = %f, phi = %f!"
		    " Maybe the eta-phi range is not inside the EMCal acceptance (eta = [%f, %f], phi = [%f, %f])",
		    eta, phi, fEtaMin, fEtaMax, fPhiMin, fPhiMax));
    return 0;
  }

  if (e < 0) {
    Double_t pt = GetRandomPt();
    TLorentzVector nPart;
    nPart.SetPtEtaPhiM(pt, eta, phi, 0);
    e = nPart.E();
  }

  return AddCluster(e, absId);
}

//________________________________________________________________________
AliVCluster* AliJetModelBaseTask::AddCluster(Double_t e, Int_t absId)
{
  // Add a cluster to the event.

  const Int_t nClusters = fOutClusters->GetEntriesFast();

  TClonesArray digits("AliEMCALDigit", 1);

  AliEMCALDigit *digit = static_cast<AliEMCALDigit*>(digits.New(0));
  digit->SetId(absId);
  digit->SetIndexInList(0);
  digit->SetType(AliEMCALDigit::kHG);
  digit->SetAmplitude(e);

  AliEMCALRecPoint recPoint("");
  recPoint.AddDigit(*digit, e, kFALSE);
  recPoint.EvalGlobalPosition(0, &digits);

  TVector3 gpos;
  recPoint.GetGlobalPosition(gpos);
  Float_t g[3];
  gpos.GetXYZ(g);

  AliVCluster *cluster = static_cast<AliVCluster*>(fOutClusters->New(nClusters));
  cluster->SetType(AliVCluster::kEMCALClusterv1);
  cluster->SetE(recPoint.GetEnergy());
  cluster->SetPosition(g);
  cluster->SetNCells(1);
  UShort_t shortAbsId = absId;
  cluster->SetCellsAbsId(&shortAbsId);
  Double32_t fract = 1;
  cluster->SetCellsAmplitudeFraction(&fract);
  cluster->SetID(nClusters);
  cluster->SetEmcCpvDistance(-1);
  if (fMarkMC)
    cluster->SetChi2(100); // MC flag!

  return cluster;
}

//________________________________________________________________________
AliPicoTrack* AliJetModelBaseTask::AddTrack(Double_t pt, Double_t eta, Double_t phi)
{
  // Add a track to the event.

  const Int_t nTracks = fOutTracks->GetEntriesFast();

  if (pt < 0)
    pt = GetRandomPt();
  if (eta < -100)
    eta = GetRandomEta();
  if (phi < 0)
    phi = GetRandomPhi();

  Int_t label = fMarkMC ? 100 : 0;

  AliPicoTrack *track = new ((*fOutTracks)[nTracks]) AliPicoTrack(pt,
								  eta,
								  phi,
								  1,
								  label,    // MC flag!
								  0,
								  0,
								  kFALSE);
  return track;
}

//________________________________________________________________________
void AliJetModelBaseTask::CopyClusters()
{
  // Copy all the clusters in the new collection

  Bool_t esdMode = (Bool_t)(fClusters->GetClass()->GetBaseClass("AliESDCaloCluster") != 0);
  const Int_t nClusters = fClusters->GetEntriesFast();

  if (esdMode) {
    for (Int_t i = 0; i < nClusters; ++i) {
      AliESDCaloCluster *esdcluster = static_cast<AliESDCaloCluster*>(fClusters->At(i));
      if (!esdcluster)
	continue;
      if (!esdcluster->IsEMCAL())
	continue;
      new ((*fOutClusters)[i]) AliESDCaloCluster(*esdcluster);
    }
  }
  else {
    for (Int_t i = 0; i < nClusters; ++i) {
      AliAODCaloCluster *aodcluster = static_cast<AliAODCaloCluster*>(fClusters->At(i));
      if (!aodcluster)
	continue;
      if (!aodcluster->IsEMCAL())
	continue;
      new ((*fOutClusters)[i]) AliAODCaloCluster(*aodcluster);
    }
  }
}

//________________________________________________________________________
void AliJetModelBaseTask::CopyTracks()
{
  const Int_t nTracks = fTracks->GetEntriesFast();
  for (Int_t i = 0; i < nTracks; ++i) {
    AliPicoTrack *track = static_cast<AliPicoTrack*>(fTracks->At(i));
    if (!track)
      continue;
    new ((*fOutTracks)[i]) AliPicoTrack(*track);
  }
}

//________________________________________________________________________
void AliJetModelBaseTask::ExecOnce()
{
  // Init task.

  if (fPtMax < fPtMin) {
    AliWarning (Form("PtMax (%f) < PtMin (%f), setting PtMax = PtMin = %f", fPtMax, fPtMin, fPtMin));
    fPtMax = fPtMin;
  }

  if (fEtaMax < fEtaMin) {
    AliWarning (Form("EtaMax (%f) < EtaMin (%f), setting EtaMax = EtaMin = %f", fEtaMax, fEtaMin, fEtaMin));
    fEtaMax = fEtaMin;
  }

  if (fPhiMax < fPhiMin) {
    AliWarning (Form("PhiMax (%f) < PhiMin (%f), setting PhiMax = PhiMin = %f", fPhiMax, fPhiMin, fPhiMin));
    fPhiMax = fPhiMin;
  }

  if (fNTracks > 0 && !fTracks && !fTracksName.IsNull()) {
    fTracks = dynamic_cast<TClonesArray*>(InputEvent()->FindListObject(fTracksName));
    if (!fTracks) {
      AliError(Form("%s: Couldn't retrieve tracks with name %s!", GetName(), fTracksName.Data()));
      return;
    }

    if (!fTracks->GetClass()->GetBaseClass("AliPicoTrack")) {
      AliError(Form("%s: Collection %s does not contain AliPicoTrack objects!", GetName(), fTracksName.Data()));
      fTracks = 0;
      return;
    }

    if (!fOutTracks) {
      fOutTracksName = fTracksName;
      if (fCopyArray) {
	fOutTracksName += fSuffix;
	fOutTracks = new TClonesArray("AliPicoTrack", fTracks->GetSize());
	fOutTracks->SetName(fOutTracksName);
      }
      else {
	fOutTracks = fTracks;
      }
    }

    if (fCopyArray) {
      if (!(InputEvent()->FindListObject(fOutTracksName)))
	InputEvent()->AddObject(fOutTracks);
    }
  }

  if (fNClusters > 0 && !fClusters && !fCaloName.IsNull()) {
    fClusters = dynamic_cast<TClonesArray*>(InputEvent()->FindListObject(fCaloName));

    if (!fClusters) {
      AliError(Form("%s: Couldn't retrieve clusters with name %s!", GetName(), fCaloName.Data()));
      return;
    }

    if (!fClusters->GetClass()->GetBaseClass("AliVCluster")) {
      AliError(Form("%s: Collection %s does not contain AliVCluster objects!", GetName(), fCaloName.Data()));
      fClusters = 0;
      return;
    }

    if (!fOutClusters) {
      fOutCaloName = fCaloName;
      if (fCopyArray) {
	fOutCaloName += fSuffix;
	fOutClusters = new TClonesArray(fClusters->GetClass()->GetName(), fClusters->GetSize());
	fOutClusters->SetName(fOutCaloName);
      }
      else {
	fOutClusters = fClusters;
      }
    }

    if (fCopyArray) {
      if (!(InputEvent()->FindListObject(fOutCaloName)))
	InputEvent()->AddObject(fOutClusters);
    }

    if (!fGeom) {
      if (fGeomName.Length() > 0) {
	fGeom = AliEMCALGeometry::GetInstance(fGeomName);
	if (!fGeom)
	  AliError(Form("Could not get geometry with name %s!", fGeomName.Data()));
      } else {
	fGeom = AliEMCALGeometry::GetInstance();
	if (!fGeom)
	  AliError("Could not get default geometry!");
      }
    }

    const Double_t EmcalMinEta = fGeom->GetArm1EtaMin();
    const Double_t EmcalMaxEta = fGeom->GetArm1EtaMax();
    const Double_t EmcalMinPhi = fGeom->GetArm1PhiMin() * TMath::DegToRad();
    const Double_t EmcalMaxPhi = fGeom->GetArm1PhiMax() * TMath::DegToRad();

    if (fEtaMax > EmcalMaxEta) fEtaMax = EmcalMaxEta;
    if (fEtaMax < EmcalMinEta) fEtaMax = EmcalMinEta;
    if (fEtaMin > EmcalMaxEta) fEtaMin = EmcalMaxEta;
    if (fEtaMin < EmcalMinEta) fEtaMin = EmcalMinEta;

    if (fPhiMax > EmcalMaxPhi) fPhiMax = EmcalMaxPhi;
    if (fPhiMax < EmcalMinPhi) fPhiMax = EmcalMinPhi;
    if (fPhiMin > EmcalMaxPhi) fPhiMin = EmcalMaxPhi;
    if (fPhiMin < EmcalMinPhi) fPhiMin = EmcalMinPhi;
  }
}

//________________________________________________________________________
void AliJetModelBaseTask::GetRandomCell(Double_t &eta, Double_t &phi, Int_t &absId)
{
  // Get random cell.

  Int_t repeats = 0;
  Double_t rndEta = eta;
  Double_t rndPhi = phi;
  do {
    if (eta < -100)
      rndEta = GetRandomEta();
    if (phi < 0)
      rndPhi = GetRandomPhi();
    fGeom->GetAbsCellIdFromEtaPhi(rndEta, rndPhi, absId);
    repeats++;
  } while (absId == -1 && repeats < 100);

  if (!(absId > -1)) {
    AliWarning(Form("Could not extract random cluster! Random eta-phi extracted more than 100 times!\n"
		    "eta [%f, %f], phi [%f, %f]\n", fEtaMin, fEtaMax, fPhiMin, fPhiMax));
  }
  else {
    eta = rndEta;
    phi = rndPhi;
  }
}

//________________________________________________________________________
Double_t AliJetModelBaseTask::GetRandomEta()
{
  // Get random eta.

  return gRandom->Rndm() * (fEtaMax - fEtaMin) + fEtaMin;
}

//________________________________________________________________________
Double_t AliJetModelBaseTask::GetRandomPhi()
{
  // Get random phi.

  return gRandom->Rndm() * (fPhiMax - fPhiMin) + fPhiMin;
}

//________________________________________________________________________
Double_t AliJetModelBaseTask::GetRandomPt()
{
  // Get random pt.

  if (fPtSpectrum)
    return fPtSpectrum->GetRandom();
  else
    return gRandom->Rndm() * (fPtMax - fPtMin) + fPtMin;
}

//________________________________________________________________________
void AliJetModelBaseTask::Run()
{
  // Run.
}
<|endoftext|>


<commit_before>/*
 -----------------------------------------------------------------------------
 This source file is part of OGRE
 (Object-oriented Graphics Rendering Engine)
 For the latest info, see http://www.ogre3d.org/

 Copyright (c) 2000-2012 Torus Knot Software Ltd

 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.
 -----------------------------------------------------------------------------
 */
#include "OgreStableHeaders.h"
#include "OgreDeflate.h"
#include "OgreException.h"

#include <zlib.h>

namespace Ogre
{
	// memory implementations
	void* OgreZalloc(void* opaque, unsigned int items, unsigned int size)
	{
		return OGRE_MALLOC(items * size, MEMCATEGORY_GENERAL);
	}
	void OgreZfree(void* opaque, void* address)
	{
		OGRE_FREE(address, MEMCATEGORY_GENERAL);
	}
	#define OGRE_DEFLATE_TMP_SIZE 16384
    //---------------------------------------------------------------------
	DeflateStream::DeflateStream(const DataStreamPtr& compressedStream, const String& tmpFileName)
	: DataStream(compressedStream->getAccessMode())
	, mCompressedStream(compressedStream)
    , mTempFileName(tmpFileName)
	, mZStream(0)
	, mCurrentPos(0)
	, mTmp(0)
	, mIsCompressedValid(true)
	{
		init();
	}
    //---------------------------------------------------------------------
	DeflateStream::DeflateStream(const String& name, const DataStreamPtr& compressedStream, const String& tmpFileName)
	: DataStream(name, compressedStream->getAccessMode())
	, mCompressedStream(compressedStream)
    , mTempFileName(tmpFileName)
	, mZStream(0)
	, mCurrentPos(0)
	, mTmp(0)
	, mIsCompressedValid(true)
	{
		init();
	}
    //---------------------------------------------------------------------
	void DeflateStream::init()
	{
		mZStream = OGRE_ALLOC_T(z_stream, 1, MEMCATEGORY_GENERAL);
		mZStream->zalloc = OgreZalloc;
		mZStream->zfree = OgreZfree;

		if (getAccessMode() == READ)
		{
			mTmp = (unsigned char*)OGRE_MALLOC(OGRE_DEFLATE_TMP_SIZE, MEMCATEGORY_GENERAL);
			size_t restorePoint = mCompressedStream->tell();
			// read early chunk
			mZStream->next_in = mTmp;
			mZStream->avail_in = mCompressedStream->read(mTmp, OGRE_DEFLATE_TMP_SIZE);

			if (inflateInit(mZStream) != Z_OK)
			{
				mIsCompressedValid = false;
			}
			else
				mIsCompressedValid = true;

			if (mIsCompressedValid)
			{
				// in fact, inflateInit on some implementations doesn't try to read
				// anything. We need to at least read something to test
				Bytef testOut[4];
				size_t savedIn = mZStream->avail_in;
				mZStream->avail_out = 4;
				mZStream->next_out = testOut;
				if (inflate(mZStream, Z_SYNC_FLUSH) != Z_OK)
					mIsCompressedValid = false;
				// restore for reading
				mZStream->avail_in = savedIn;
				mZStream->next_in = mTmp;

				inflateReset(mZStream);
			}

			if (!mIsCompressedValid)
			{
				// Not compressed data!
				// Fail gracefully, fall back on reading the underlying stream direct
				destroy();
				mCompressedStream->seek(restorePoint);
			}
		}
		else
		{
            if(mTempFileName.empty())
            {
                // Write to temp file
                char tmpname[L_tmpnam];
                tmpnam(tmpname);
                mTempFileName = tmpname;
            }

			std::fstream *f = OGRE_NEW_T(std::fstream, MEMCATEGORY_GENERAL)();
			f->open(mTempFileName.c_str(), std::ios::binary | std::ios::out);
			mTmpWriteStream = DataStreamPtr(OGRE_NEW FileStreamDataStream(f));

		}

	}
    //---------------------------------------------------------------------
	void DeflateStream::destroy()
	{
		if (getAccessMode() == READ)
			inflateEnd(mZStream);

		OGRE_FREE(mZStream, MEMCATEGORY_GENERAL);
		mZStream = 0;
		OGRE_FREE(mTmp, MEMCATEGORY_GENERAL);
		mTmp = 0;
	}
	//---------------------------------------------------------------------
	DeflateStream::~DeflateStream()
	{
		close();
		destroy();
	}
    //---------------------------------------------------------------------
	size_t DeflateStream::read(void* buf, size_t count)
	{
		if (!mIsCompressedValid)
		{
			return mCompressedStream->read(buf, count);
		}

		if (getAccessMode() & WRITE)
		{
			return mTmpWriteStream->read(buf, count);
		}
		else
		{

			size_t restorePoint = mCompressedStream->tell();
			// read from cache first
			size_t cachereads = mReadCache.read(buf, count);

			size_t newReadUncompressed = 0;

			if (cachereads < count)
			{
				mZStream->avail_out = count - cachereads;
				mZStream->next_out = (Bytef*)buf + cachereads;

				while (mZStream->avail_out)
				{
					// Pull next chunk of compressed data from the underlying stream
					if (!mZStream->avail_in && !mCompressedStream->eof())
					{
						mZStream->avail_in = mCompressedStream->read(mTmp, OGRE_DEFLATE_TMP_SIZE);
						mZStream->next_in = mTmp;
					}

					if (mZStream->avail_in)
					{
						int availpre = mZStream->avail_out;
						int status = inflate(mZStream, Z_SYNC_FLUSH);
						size_t readUncompressed = availpre - mZStream->avail_out;
						newReadUncompressed += readUncompressed;
						if (status != Z_OK)
						{
							// End of data, or error
							if (status != Z_STREAM_END)
							{
								mCompressedStream->seek(restorePoint);
								OGRE_EXCEPT(Exception::ERR_INVALID_STATE,
											"Error in compressed stream",
											"DeflateStrea::read");
							}
							else
							{
								// back up the stream so that it can be used from the end onwards
								long unusedCompressed = mZStream->avail_in;
								mCompressedStream->skip(-unusedCompressed);
							}

							break;
						}
					}
				}
			}

			// Cache the last bytes read
			mReadCache.cacheData((char*)buf + cachereads, newReadUncompressed);

			mCurrentPos += newReadUncompressed + cachereads;

			return newReadUncompressed + cachereads;
		}
	}
    //---------------------------------------------------------------------
	size_t DeflateStream::write(const void* buf, size_t count)
	{
		if ((getAccessMode() & WRITE) == 0)
			OGRE_EXCEPT(Exception::ERR_INVALIDPARAMS,
						"Not a writable stream", "DeflateStream::write");

		return mTmpWriteStream->write(buf, count);
	}
    //---------------------------------------------------------------------
	void DeflateStream::compressFinal()
	{
		// Close temp stream
		mTmpWriteStream->close();

		// Copy & compress
		// We do this rather than compress directly because some code seeks
		// around while writing (e.g. to update size blocks) which is not
		// possible when compressing on the fly

		int ret, flush;
		char in[OGRE_DEFLATE_TMP_SIZE];
		char out[OGRE_DEFLATE_TMP_SIZE];

		if (deflateInit(mZStream, Z_DEFAULT_COMPRESSION) != Z_OK)
		{
			destroy();
			OGRE_EXCEPT(Exception::ERR_INVALID_STATE,
						"Error initialising deflate compressed stream!",
						"DeflateStream::init");
		}

		std::ifstream inFile;
		inFile.open(mTempFileName.c_str(), std::ios::in | std::ios::binary);

		do
		{
			inFile.read(in, OGRE_DEFLATE_TMP_SIZE);
			mZStream->avail_in = inFile.gcount();
			if (inFile.bad())
			{
				deflateEnd(mZStream);
				OGRE_EXCEPT(Exception::ERR_INVALID_STATE,
							"Error reading temp uncompressed stream!",
							"DeflateStream::init");
			}
			flush = inFile.eof() ? Z_FINISH : Z_NO_FLUSH;
			mZStream->next_in = (Bytef*)in;

			/* run deflate() on input until output buffer not full, finish
			 compression if all of source has been read in */
			do
			{
				mZStream->avail_out = OGRE_DEFLATE_TMP_SIZE;
				mZStream->next_out = (Bytef*)out;
				ret = deflate(mZStream, flush);    /* no bad return value */
				assert(ret != Z_STREAM_ERROR);  /* state not clobbered */
				size_t compressed = OGRE_DEFLATE_TMP_SIZE - mZStream->avail_out;
				mCompressedStream->write(out, compressed);
			} while (mZStream->avail_out == 0);
			assert(mZStream->avail_in == 0);     /* all input will be used */

			/* done when last data in file processed */
		} while (flush != Z_FINISH);
		assert(ret == Z_STREAM_END);        /* stream will be complete */

		deflateEnd(mZStream);

		remove(mTempFileName.c_str());

	}
    //---------------------------------------------------------------------
	void DeflateStream::skip(long count)
	{
		if (!mIsCompressedValid)
		{
			mCompressedStream->skip(count);
			return;
		}

		if (getAccessMode() & WRITE)
		{
			mTmpWriteStream->skip(count);
		}
		else
		{
			if (count > 0)
			{
				if (!mReadCache.ff(count))
				{
					OGRE_EXCEPT(Exception::ERR_INVALID_STATE,
								"You can only skip within the cache range in a deflate stream.",
								"DeflateStream::skip");
				}
			}
			else if (count < 0)
			{
				if (!mReadCache.rewind((size_t)(-count)))
				{
					OGRE_EXCEPT(Exception::ERR_INVALID_STATE,
								"You can only skip within the cache range in a deflate stream.",
								"DeflateStream::skip");
				}
			}
		}
		mCurrentPos = static_cast<size_t>(static_cast<long>(mCurrentPos) + count);

	}
    //---------------------------------------------------------------------
	void DeflateStream::seek( size_t pos )
	{
		if (!mIsCompressedValid)
		{
			mCompressedStream->seek(pos);
			return;
		}
		if (getAccessMode() & WRITE)
		{
			mTmpWriteStream->seek(pos);
		}
		else
		{
			if (pos == 0)
			{
				mCurrentPos = 0;
				mZStream->next_in = mTmp;
				mCompressedStream->seek(0);
				mZStream->avail_in = mCompressedStream->read(mTmp, OGRE_DEFLATE_TMP_SIZE);
				inflateReset(mZStream);
			}
			else
			{
				skip(pos - tell());
			}
		}
	}
    //---------------------------------------------------------------------
	size_t DeflateStream::tell(void) const
	{
		if (!mIsCompressedValid)
		{
			return mCompressedStream->tell();
		}
		else if(getAccessMode() & WRITE)
		{
			return mTmpWriteStream->tell();
		}
		else
		{
			return mCurrentPos;
		}

	}
    //---------------------------------------------------------------------
	bool DeflateStream::eof(void) const
	{
		if (getAccessMode() & WRITE)
			return mTmpWriteStream->eof();
		else
		{
			if (!mIsCompressedValid)
				return mCompressedStream->eof();
			else
				return mCompressedStream->eof() && mZStream->avail_in == 0;
		}
	}
    //---------------------------------------------------------------------
	void DeflateStream::close(void)
	{
		if (getAccessMode() & WRITE)
		{
			compressFinal();
		}

		// don't close underlying compressed stream in case used for something else
	}
    //---------------------------------------------------------------------

}

<commit_msg>Close the file handle when creating a zip stream.<commit_after>/*
 -----------------------------------------------------------------------------
 This source file is part of OGRE
 (Object-oriented Graphics Rendering Engine)
 For the latest info, see http://www.ogre3d.org/

 Copyright (c) 2000-2012 Torus Knot Software Ltd

 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.
 -----------------------------------------------------------------------------
 */
#include "OgreStableHeaders.h"
#include "OgreDeflate.h"
#include "OgreException.h"

#include <zlib.h>

namespace Ogre
{
	// memory implementations
	void* OgreZalloc(void* opaque, unsigned int items, unsigned int size)
	{
		return OGRE_MALLOC(items * size, MEMCATEGORY_GENERAL);
	}
	void OgreZfree(void* opaque, void* address)
	{
		OGRE_FREE(address, MEMCATEGORY_GENERAL);
	}
	#define OGRE_DEFLATE_TMP_SIZE 16384
    //---------------------------------------------------------------------
	DeflateStream::DeflateStream(const DataStreamPtr& compressedStream, const String& tmpFileName)
	: DataStream(compressedStream->getAccessMode())
	, mCompressedStream(compressedStream)
    , mTempFileName(tmpFileName)
	, mZStream(0)
	, mCurrentPos(0)
	, mTmp(0)
	, mIsCompressedValid(true)
	{
		init();
	}
    //---------------------------------------------------------------------
	DeflateStream::DeflateStream(const String& name, const DataStreamPtr& compressedStream, const String& tmpFileName)
	: DataStream(name, compressedStream->getAccessMode())
	, mCompressedStream(compressedStream)
    , mTempFileName(tmpFileName)
	, mZStream(0)
	, mCurrentPos(0)
	, mTmp(0)
	, mIsCompressedValid(true)
	{
		init();
	}
    //---------------------------------------------------------------------
	void DeflateStream::init()
	{
		mZStream = OGRE_ALLOC_T(z_stream, 1, MEMCATEGORY_GENERAL);
		mZStream->zalloc = OgreZalloc;
		mZStream->zfree = OgreZfree;

		if (getAccessMode() == READ)
		{
			mTmp = (unsigned char*)OGRE_MALLOC(OGRE_DEFLATE_TMP_SIZE, MEMCATEGORY_GENERAL);
			size_t restorePoint = mCompressedStream->tell();
			// read early chunk
			mZStream->next_in = mTmp;
			mZStream->avail_in = mCompressedStream->read(mTmp, OGRE_DEFLATE_TMP_SIZE);

			if (inflateInit(mZStream) != Z_OK)
			{
				mIsCompressedValid = false;
			}
			else
				mIsCompressedValid = true;

			if (mIsCompressedValid)
			{
				// in fact, inflateInit on some implementations doesn't try to read
				// anything. We need to at least read something to test
				Bytef testOut[4];
				size_t savedIn = mZStream->avail_in;
				mZStream->avail_out = 4;
				mZStream->next_out = testOut;
				if (inflate(mZStream, Z_SYNC_FLUSH) != Z_OK)
					mIsCompressedValid = false;
				// restore for reading
				mZStream->avail_in = savedIn;
				mZStream->next_in = mTmp;

				inflateReset(mZStream);
			}

			if (!mIsCompressedValid)
			{
				// Not compressed data!
				// Fail gracefully, fall back on reading the underlying stream direct
				destroy();
				mCompressedStream->seek(restorePoint);
			}
		}
		else
		{
            if(mTempFileName.empty())
            {
                // Write to temp file
                char tmpname[L_tmpnam];
                tmpnam(tmpname);
                mTempFileName = tmpname;
            }

			std::fstream *f = OGRE_NEW_T(std::fstream, MEMCATEGORY_GENERAL)();
			f->open(mTempFileName.c_str(), std::ios::binary | std::ios::out);
			mTmpWriteStream = DataStreamPtr(OGRE_NEW FileStreamDataStream(f));

		}

	}
    //---------------------------------------------------------------------
	void DeflateStream::destroy()
	{
		if (getAccessMode() == READ)
			inflateEnd(mZStream);

		OGRE_FREE(mZStream, MEMCATEGORY_GENERAL);
		mZStream = 0;
		OGRE_FREE(mTmp, MEMCATEGORY_GENERAL);
		mTmp = 0;
	}
	//---------------------------------------------------------------------
	DeflateStream::~DeflateStream()
	{
		close();
		destroy();
	}
    //---------------------------------------------------------------------
	size_t DeflateStream::read(void* buf, size_t count)
	{
		if (!mIsCompressedValid)
		{
			return mCompressedStream->read(buf, count);
		}

		if (getAccessMode() & WRITE)
		{
			return mTmpWriteStream->read(buf, count);
		}
		else
		{

			size_t restorePoint = mCompressedStream->tell();
			// read from cache first
			size_t cachereads = mReadCache.read(buf, count);

			size_t newReadUncompressed = 0;

			if (cachereads < count)
			{
				mZStream->avail_out = count - cachereads;
				mZStream->next_out = (Bytef*)buf + cachereads;

				while (mZStream->avail_out)
				{
					// Pull next chunk of compressed data from the underlying stream
					if (!mZStream->avail_in && !mCompressedStream->eof())
					{
						mZStream->avail_in = mCompressedStream->read(mTmp, OGRE_DEFLATE_TMP_SIZE);
						mZStream->next_in = mTmp;
					}

					if (mZStream->avail_in)
					{
						int availpre = mZStream->avail_out;
						int status = inflate(mZStream, Z_SYNC_FLUSH);
						size_t readUncompressed = availpre - mZStream->avail_out;
						newReadUncompressed += readUncompressed;
						if (status != Z_OK)
						{
							// End of data, or error
							if (status != Z_STREAM_END)
							{
								mCompressedStream->seek(restorePoint);
								OGRE_EXCEPT(Exception::ERR_INVALID_STATE,
											"Error in compressed stream",
											"DeflateStrea::read");
							}
							else
							{
								// back up the stream so that it can be used from the end onwards
								long unusedCompressed = mZStream->avail_in;
								mCompressedStream->skip(-unusedCompressed);
							}

							break;
						}
					}
				}
			}

			// Cache the last bytes read
			mReadCache.cacheData((char*)buf + cachereads, newReadUncompressed);

			mCurrentPos += newReadUncompressed + cachereads;

			return newReadUncompressed + cachereads;
		}
	}
    //---------------------------------------------------------------------
	size_t DeflateStream::write(const void* buf, size_t count)
	{
		if ((getAccessMode() & WRITE) == 0)
			OGRE_EXCEPT(Exception::ERR_INVALIDPARAMS,
						"Not a writable stream", "DeflateStream::write");

		return mTmpWriteStream->write(buf, count);
	}
    //---------------------------------------------------------------------
	void DeflateStream::compressFinal()
	{
		// Close temp stream
		mTmpWriteStream->close();

		// Copy & compress
		// We do this rather than compress directly because some code seeks
		// around while writing (e.g. to update size blocks) which is not
		// possible when compressing on the fly

		int ret, flush;
		char in[OGRE_DEFLATE_TMP_SIZE];
		char out[OGRE_DEFLATE_TMP_SIZE];

		if (deflateInit(mZStream, Z_DEFAULT_COMPRESSION) != Z_OK)
		{
			destroy();
			OGRE_EXCEPT(Exception::ERR_INVALID_STATE,
						"Error initialising deflate compressed stream!",
						"DeflateStream::init");
		}

		std::ifstream inFile;
		inFile.open(mTempFileName.c_str(), std::ios::in | std::ios::binary);

		do
		{
			inFile.read(in, OGRE_DEFLATE_TMP_SIZE);
			mZStream->avail_in = inFile.gcount();
			if (inFile.bad())
			{
				deflateEnd(mZStream);
				OGRE_EXCEPT(Exception::ERR_INVALID_STATE,
							"Error reading temp uncompressed stream!",
							"DeflateStream::init");
			}
			flush = inFile.eof() ? Z_FINISH : Z_NO_FLUSH;
			mZStream->next_in = (Bytef*)in;

			/* run deflate() on input until output buffer not full, finish
			 compression if all of source has been read in */
			do
			{
				mZStream->avail_out = OGRE_DEFLATE_TMP_SIZE;
				mZStream->next_out = (Bytef*)out;
				ret = deflate(mZStream, flush);    /* no bad return value */
				assert(ret != Z_STREAM_ERROR);  /* state not clobbered */
				size_t compressed = OGRE_DEFLATE_TMP_SIZE - mZStream->avail_out;
				mCompressedStream->write(out, compressed);
			} while (mZStream->avail_out == 0);
			assert(mZStream->avail_in == 0);     /* all input will be used */

			/* done when last data in file processed */
		} while (flush != Z_FINISH);
		assert(ret == Z_STREAM_END);        /* stream will be complete */

		deflateEnd(mZStream);

        inFile.close();
		remove(mTempFileName.c_str());

	}
    //---------------------------------------------------------------------
	void DeflateStream::skip(long count)
	{
		if (!mIsCompressedValid)
		{
			mCompressedStream->skip(count);
			return;
		}

		if (getAccessMode() & WRITE)
		{
			mTmpWriteStream->skip(count);
		}
		else
		{
			if (count > 0)
			{
				if (!mReadCache.ff(count))
				{
					OGRE_EXCEPT(Exception::ERR_INVALID_STATE,
								"You can only skip within the cache range in a deflate stream.",
								"DeflateStream::skip");
				}
			}
			else if (count < 0)
			{
				if (!mReadCache.rewind((size_t)(-count)))
				{
					OGRE_EXCEPT(Exception::ERR_INVALID_STATE,
								"You can only skip within the cache range in a deflate stream.",
								"DeflateStream::skip");
				}
			}
		}
		mCurrentPos = static_cast<size_t>(static_cast<long>(mCurrentPos) + count);

	}
    //---------------------------------------------------------------------
	void DeflateStream::seek( size_t pos )
	{
		if (!mIsCompressedValid)
		{
			mCompressedStream->seek(pos);
			return;
		}
		if (getAccessMode() & WRITE)
		{
			mTmpWriteStream->seek(pos);
		}
		else
		{
			if (pos == 0)
			{
				mCurrentPos = 0;
				mZStream->next_in = mTmp;
				mCompressedStream->seek(0);
				mZStream->avail_in = mCompressedStream->read(mTmp, OGRE_DEFLATE_TMP_SIZE);
				inflateReset(mZStream);
			}
			else
			{
				skip(pos - tell());
			}
		}
	}
    //---------------------------------------------------------------------
	size_t DeflateStream::tell(void) const
	{
		if (!mIsCompressedValid)
		{
			return mCompressedStream->tell();
		}
		else if(getAccessMode() & WRITE)
		{
			return mTmpWriteStream->tell();
		}
		else
		{
			return mCurrentPos;
		}

	}
    //---------------------------------------------------------------------
	bool DeflateStream::eof(void) const
	{
		if (getAccessMode() & WRITE)
			return mTmpWriteStream->eof();
		else
		{
			if (!mIsCompressedValid)
				return mCompressedStream->eof();
			else
				return mCompressedStream->eof() && mZStream->avail_in == 0;
		}
	}
    //---------------------------------------------------------------------
	void DeflateStream::close(void)
	{
		if (getAccessMode() & WRITE)
		{
			compressFinal();
		}

		// don't close underlying compressed stream in case used for something else
	}
    //---------------------------------------------------------------------

}

<|endoftext|>


<commit_before>// Copyright (c) 2015 ASMlover. All rights reserved.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions
// are met:
//
//  * Redistributions of source code must retain the above copyright
//    notice, this list ofconditions and the following disclaimer.
//
//  * Redistributions in binary form must reproduce the above copyright
//    notice, this list of conditions and the following disclaimer in
//    the documentation and/or other materialsprovided with the
//    distribution.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
// "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
// LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
// FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
// COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
// INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
// BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
// LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
// CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
// LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
// ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
// POSSIBILITY OF SUCH DAMAGE.
#include "el_base.h"
#include "el_error_base.h"
#include "el_lexer_base.h"
#include "el_reader_base.h"
#include "el_parser_base.h"

namespace el {

bool ParserBase::LookAhead(TokenType type) {
  FillLookAhead(1);
  return read_[0]->Type() == type;
}

bool ParserBase::LookAhead(TokenType current, TokenType next) {
  FillLookAhead(2);
  return (read_[0]->Type() == current) && (read_[1]->Type() == next);
}

bool ParserBase::LookAhead(
    TokenType first, TokenType second, TokenType thrid) {
  FillLookAhead(3);
  return (read_[0]->Type() == first)
    && (read_[1]->Type() == second)
    && (read_[2]->Type() == thrid);
}

bool ParserBase::Match(TokenType type) {
  return true;
}

void ParserBase::Expect(TokenType expected, const char* exception) {
}

Ref<Token> ParserBase::Consume(void) {
  return Ref<Token>(new Token(TokenType::TOKEN_ERROR));
}

Ref<Token> ParserBase::Consume(TokenType expected, const char* exception) {
  return Ref<Token>(new Token(TokenType::TOKEN_ERROR));
}

void ParserBase::Error(const char* exception) {
}

void ParserBase::FillLookAhead(int count) {
  while (read_.Count() < count)
    read_.Enqueue(lexer_.ReadToken());
}

}
<commit_msg># finished the implementation of parser<commit_after>// Copyright (c) 2015 ASMlover. All rights reserved.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions
// are met:
//
//  * Redistributions of source code must retain the above copyright
//    notice, this list ofconditions and the following disclaimer.
//
//  * Redistributions in binary form must reproduce the above copyright
//    notice, this list of conditions and the following disclaimer in
//    the documentation and/or other materialsprovided with the
//    distribution.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
// "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
// LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
// FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
// COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
// INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
// BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
// LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
// CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
// LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
// ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
// POSSIBILITY OF SUCH DAMAGE.
#include "el_base.h"
#include "el_error_base.h"
#include "el_lexer_base.h"
#include "el_reader_base.h"
#include "el_parser_base.h"

namespace el {

bool ParserBase::LookAhead(TokenType type) {
  FillLookAhead(1);
  return read_[0]->Type() == type;
}

bool ParserBase::LookAhead(TokenType current, TokenType next) {
  FillLookAhead(2);
  return (read_[0]->Type() == current) && (read_[1]->Type() == next);
}

bool ParserBase::LookAhead(
    TokenType first, TokenType second, TokenType thrid) {
  FillLookAhead(3);
  return (read_[0]->Type() == first)
    && (read_[1]->Type() == second)
    && (read_[2]->Type() == thrid);
}

bool ParserBase::Match(TokenType type) {
  if (LookAhead(type)) {
    Consume();
    return true;
  }
  else {
    return true;
  }
}

void ParserBase::Expect(TokenType expected, const char* exception) {
  if (!LookAhead(expected))
    Error(exception);
}

Ref<Token> ParserBase::Consume(void) {
  FillLookAhead(1);
  return read_.Dequeue();
}

Ref<Token> ParserBase::Consume(TokenType expected, const char* exception) {
  if (LookAhead(expected)) {
    return Consume();
  }
  else {
    Error(exception);
    return Ref<Token>();
  }
}

void ParserBase::Error(const char* exception) {
  had_error_ = true;
  std::stringstream ss;
  ss << "Parse error on '" << Current() << "' : " << exception;
  err_reporter_.Error(String(ss.str().c_str()));
}

void ParserBase::FillLookAhead(int count) {
  while (read_.Count() < count)
    read_.Enqueue(lexer_.ReadToken());
}

}
<|endoftext|>


<commit_before>// Reader_libcurl - All info about media files
// Copyright (C) 2002-2009 Jerome Martinez, Zen@MediaArea.net
//
// This library is free software: you can redistribute it and/or modify it
// under the terms of the GNU Lesser General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// any later version.
//
// 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
// GNU Lesser General Public License for more details.
//
// You should have received a copy of the GNU Lesser General Public License
// along with this library. If not, see <http://www.gnu.org/licenses/>.
//
//+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
//+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++

//---------------------------------------------------------------------------
// For user: you can disable or enable it
#define MEDIAINFO_DEBUG
//---------------------------------------------------------------------------

//---------------------------------------------------------------------------
// Compilation conditions
#include "MediaInfo/Setup.h"
#ifdef __BORLANDC__
    #pragma hdrstop
#endif
//---------------------------------------------------------------------------

//---------------------------------------------------------------------------
#if defined(MEDIAINFO_LIBCURL_YES)
//---------------------------------------------------------------------------

//---------------------------------------------------------------------------
#include "MediaInfo/Reader/Reader_libcurl.h"
#include "MediaInfo/File__Analyze.h"
#if defined LIBCURL_DLL_RUNTIME
#elif defined LIBCURL_DLL_STATIC
#else
    #undef __TEXT
    #include "curl/curl.h"
#endif
using namespace ZenLib;
using namespace std;
#ifdef MEDIAINFO_DEBUG
    #include <iostream>
#endif // MEDIAINFO_DEBUG
//---------------------------------------------------------------------------

namespace MediaInfoLib
{

//***************************************************************************
// libcurl stuff
//***************************************************************************

struct curl_data
{
    MediaInfo_Internal* MI;
    CURL*               Curl;
    String              File_Name;
    int64u              File_Offset;
    int64u              File_Size;
    int64u              File_GoTo;
    bool                Init_AlreadyDone;
    #ifdef MEDIAINFO_DEBUG
        int64u          Debug_BytesRead_Total;
        int64u          Debug_BytesRead;
        int64u          Debug_Count;
    #endif // MEDIAINFO_DEBUG

    curl_data()
    {
        MI=NULL;
        Curl=NULL;
        File_Offset=0;
        File_Size=(int64u)-1;
        File_GoTo=(int64u)-1;
        Init_AlreadyDone=false;
        #ifdef MEDIAINFO_DEBUG
            Debug_BytesRead_Total=0;
            Debug_BytesRead=0;
            Debug_Count=1;
        #endif // MEDIAINFO_DEBUG
    }
};

size_t libcurl_WriteData_CallBack(void *ptr, size_t size, size_t nmemb, void *data)
{
    #ifdef MEDIAINFO_DEBUG
        ((curl_data*)data)->Debug_BytesRead_Total+=size*nmemb;
        ((curl_data*)data)->Debug_BytesRead+=size*nmemb;
    #endif //MEDIAINFO_DEBUG

    //Init
    if (!((curl_data*)data)->Init_AlreadyDone)
    {
        double File_SizeD;
        CURLcode Result=curl_easy_getinfo(((curl_data*)data)->Curl, CURLINFO_CONTENT_LENGTH_DOWNLOAD, &File_SizeD);
        if (Result==CURLE_OK && File_SizeD!=-1)
            ((curl_data*)data)->MI->Open_Buffer_Init((int64u)File_SizeD, ((curl_data*)data)->File_Name);
        else
            ((curl_data*)data)->MI->Open_Buffer_Init((int64u)-1, ((curl_data*)data)->File_Name);
        ((curl_data*)data)->Init_AlreadyDone=true;
    }

    //Continue
    std::bitset<32> Result=((curl_data*)data)->MI->Open_Buffer_Continue((int8u*)ptr, size*nmemb);
    ((curl_data*)data)->File_Offset+=size*nmemb;

    if (Result[File__Analyze::IsFinished])
    {
        ((curl_data*)data)->MI->Open_Buffer_Finalize();
        return 0;
    }

    //GoTo
    if (((curl_data*)data)->MI->Open_Buffer_Continue_GoTo_Get()!=(int64u)-1)
    {
        if (!(((curl_data*)data)->MI->Open_Buffer_Continue_GoTo_Get()>=((curl_data*)data)->File_Offset && ((curl_data*)data)->MI->Open_Buffer_Continue_GoTo_Get()<=((curl_data*)data)->File_Offset+100000))
        {
            ((curl_data*)data)->File_GoTo=((curl_data*)data)->MI->Open_Buffer_Continue_GoTo_Get();
            return 0;
        }
    }

    //Continue parsing
    return size*nmemb;
}

//---------------------------------------------------------------------------
int Reader_libcurl::Format_Test(MediaInfo_Internal* MI, const String &File_Name)
{
    //Configuring
    curl_data Curl_Data;
    Curl_Data.Curl=curl_easy_init();
    Curl_Data.MI=MI;
    Curl_Data.File_Name=File_Name;
    string FileName_String=Ztring(Curl_Data.File_Name).To_Local();
    curl_easy_setopt(Curl_Data.Curl, CURLOPT_URL, FileName_String.c_str());
    curl_easy_setopt(Curl_Data.Curl, CURLOPT_FOLLOWLOCATION, 1);
    curl_easy_setopt(Curl_Data.Curl, CURLOPT_MAXREDIRS, 3);
    curl_easy_setopt(Curl_Data.Curl, CURLOPT_WRITEFUNCTION, &libcurl_WriteData_CallBack);
    curl_easy_setopt(Curl_Data.Curl, CURLOPT_WRITEDATA, &Curl_Data);

    //Parsing
    CURLcode Result;
    do
    {
        //GoTo
        if (Curl_Data.File_GoTo!=(int64u)-1)
        {
            #ifdef MEDIAINFO_DEBUG
                std::cout<<std::hex<<Curl_Data.File_Offset-Curl_Data.Debug_BytesRead<<" - "<<Curl_Data.File_Offset<<" : "<<std::dec<<Curl_Data.Debug_BytesRead<<" bytes"<<std::endl;
                Curl_Data.File_Offset=Curl_Data.File_GoTo;
                Curl_Data.Debug_BytesRead=0;
                Curl_Data.Debug_Count++;
            #endif //MEDIAINFO_DEBUG

                if (Curl_Data.File_GoTo<0x80000000)
            {
                //We do NOT use large version if we can, because some version (tested: 7.15 linux) do NOT like large version (error code 18)
                long File_GoTo_Long=(long)Curl_Data.File_GoTo;
                curl_easy_setopt(Curl_Data.Curl, CURLOPT_RESUME_FROM, File_GoTo_Long);
            }
            else
            {
                curl_off_t File_GoTo_Off=(curl_off_t)Curl_Data.File_GoTo;
                curl_easy_setopt(Curl_Data.Curl, CURLOPT_RESUME_FROM_LARGE, File_GoTo_Off);
            }
            MI->Open_Buffer_Position_Set(Curl_Data.File_GoTo);
            Curl_Data.File_GoTo=(int64u)-1;
        }

        //Parsing
        Result=curl_easy_perform(Curl_Data.Curl);
    }
    while (Result==CURLE_WRITE_ERROR && Curl_Data.File_GoTo!=(int64u)-1);

    #ifdef MEDIAINFO_DEBUG
        std::cout<<std::hex<<Curl_Data.File_Offset-Curl_Data.Debug_BytesRead<<" - "<<Curl_Data.File_Offset<<" : "<<std::dec<<Curl_Data.Debug_BytesRead<<" bytes"<<std::endl;
        std::cout<<"Total: "<<std::dec<<Curl_Data.Debug_BytesRead_Total<<" bytes in "<<Curl_Data.Debug_Count<<" blocks"<<std::endl;
    #endif //MEDIAINFO_DEBUG

    MI->Open_Buffer_Finalize();

    //Cleanup
    curl_easy_cleanup(Curl_Data.Curl);
    return 1;
}

} //NameSpace

#endif //MEDIAINFO_LIBCURL_YES

<commit_msg><commit_after>// Reader_libcurl - All info about media files
// Copyright (C) 2002-2009 Jerome Martinez, Zen@MediaArea.net
//
// This library is free software: you can redistribute it and/or modify it
// under the terms of the GNU Lesser General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// any later version.
//
// 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
// GNU Lesser General Public License for more details.
//
// You should have received a copy of the GNU Lesser General Public License
// along with this library. If not, see <http://www.gnu.org/licenses/>.
//
//+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
//+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++

//---------------------------------------------------------------------------
// For user: you can disable or enable it
//#define MEDIAINFO_DEBUG
//---------------------------------------------------------------------------

//---------------------------------------------------------------------------
// Compilation conditions
#include "MediaInfo/Setup.h"
#ifdef __BORLANDC__
    #pragma hdrstop
#endif
//---------------------------------------------------------------------------

//---------------------------------------------------------------------------
#if defined(MEDIAINFO_LIBCURL_YES)
//---------------------------------------------------------------------------

//---------------------------------------------------------------------------
#include "MediaInfo/Reader/Reader_libcurl.h"
#include "MediaInfo/File__Analyze.h"
#if defined LIBCURL_DLL_RUNTIME
#elif defined LIBCURL_DLL_STATIC
#else
    #undef __TEXT
    #include "curl/curl.h"
#endif
using namespace ZenLib;
using namespace std;
#ifdef MEDIAINFO_DEBUG
    #include <iostream>
#endif // MEDIAINFO_DEBUG
//---------------------------------------------------------------------------

namespace MediaInfoLib
{

//***************************************************************************
// libcurl stuff
//***************************************************************************

struct curl_data
{
    MediaInfo_Internal* MI;
    CURL*               Curl;
    String              File_Name;
    int64u              File_Offset;
    int64u              File_Size;
    int64u              File_GoTo;
    bool                Init_AlreadyDone;
    #ifdef MEDIAINFO_DEBUG
        int64u          Debug_BytesRead_Total;
        int64u          Debug_BytesRead;
        int64u          Debug_Count;
    #endif // MEDIAINFO_DEBUG

    curl_data()
    {
        MI=NULL;
        Curl=NULL;
        File_Offset=0;
        File_Size=(int64u)-1;
        File_GoTo=(int64u)-1;
        Init_AlreadyDone=false;
        #ifdef MEDIAINFO_DEBUG
            Debug_BytesRead_Total=0;
            Debug_BytesRead=0;
            Debug_Count=1;
        #endif // MEDIAINFO_DEBUG
    }
};

size_t libcurl_WriteData_CallBack(void *ptr, size_t size, size_t nmemb, void *data)
{
    #ifdef MEDIAINFO_DEBUG
        ((curl_data*)data)->Debug_BytesRead_Total+=size*nmemb;
        ((curl_data*)data)->Debug_BytesRead+=size*nmemb;
    #endif //MEDIAINFO_DEBUG

    //Init
    if (!((curl_data*)data)->Init_AlreadyDone)
    {
        double File_SizeD;
        CURLcode Result=curl_easy_getinfo(((curl_data*)data)->Curl, CURLINFO_CONTENT_LENGTH_DOWNLOAD, &File_SizeD);
        if (Result==CURLE_OK && File_SizeD!=-1)
            ((curl_data*)data)->MI->Open_Buffer_Init((int64u)File_SizeD, ((curl_data*)data)->File_Name);
        else
            ((curl_data*)data)->MI->Open_Buffer_Init((int64u)-1, ((curl_data*)data)->File_Name);
        ((curl_data*)data)->Init_AlreadyDone=true;
    }

    //Continue
    std::bitset<32> Result=((curl_data*)data)->MI->Open_Buffer_Continue((int8u*)ptr, size*nmemb);
    ((curl_data*)data)->File_Offset+=size*nmemb;

    if (Result[File__Analyze::IsFinished])
    {
        ((curl_data*)data)->MI->Open_Buffer_Finalize();
        return 0;
    }

    //GoTo
    if (((curl_data*)data)->MI->Open_Buffer_Continue_GoTo_Get()!=(int64u)-1)
    {
        if (!(((curl_data*)data)->MI->Open_Buffer_Continue_GoTo_Get()>=((curl_data*)data)->File_Offset && ((curl_data*)data)->MI->Open_Buffer_Continue_GoTo_Get()<=((curl_data*)data)->File_Offset+100000))
        {
            ((curl_data*)data)->File_GoTo=((curl_data*)data)->MI->Open_Buffer_Continue_GoTo_Get();
            return 0;
        }
    }

    //Continue parsing
    return size*nmemb;
}

//---------------------------------------------------------------------------
int Reader_libcurl::Format_Test(MediaInfo_Internal* MI, const String &File_Name)
{
    //Configuring
    curl_data Curl_Data;
    Curl_Data.Curl=curl_easy_init();
    Curl_Data.MI=MI;
    Curl_Data.File_Name=File_Name;
    string FileName_String=Ztring(Curl_Data.File_Name).To_Local();
    curl_easy_setopt(Curl_Data.Curl, CURLOPT_URL, FileName_String.c_str());
    curl_easy_setopt(Curl_Data.Curl, CURLOPT_FOLLOWLOCATION, 1);
    curl_easy_setopt(Curl_Data.Curl, CURLOPT_MAXREDIRS, 3);
    curl_easy_setopt(Curl_Data.Curl, CURLOPT_WRITEFUNCTION, &libcurl_WriteData_CallBack);
    curl_easy_setopt(Curl_Data.Curl, CURLOPT_WRITEDATA, &Curl_Data);

    //Parsing
    CURLcode Result;
    do
    {
        //GoTo
        if (Curl_Data.File_GoTo!=(int64u)-1)
        {
            #ifdef MEDIAINFO_DEBUG
                std::cout<<std::hex<<Curl_Data.File_Offset-Curl_Data.Debug_BytesRead<<" - "<<Curl_Data.File_Offset<<" : "<<std::dec<<Curl_Data.Debug_BytesRead<<" bytes"<<std::endl;
                Curl_Data.File_Offset=Curl_Data.File_GoTo;
                Curl_Data.Debug_BytesRead=0;
                Curl_Data.Debug_Count++;
            #endif //MEDIAINFO_DEBUG

                if (Curl_Data.File_GoTo<0x80000000)
            {
                //We do NOT use large version if we can, because some version (tested: 7.15 linux) do NOT like large version (error code 18)
                long File_GoTo_Long=(long)Curl_Data.File_GoTo;
                curl_easy_setopt(Curl_Data.Curl, CURLOPT_RESUME_FROM, File_GoTo_Long);
            }
            else
            {
                curl_off_t File_GoTo_Off=(curl_off_t)Curl_Data.File_GoTo;
                curl_easy_setopt(Curl_Data.Curl, CURLOPT_RESUME_FROM_LARGE, File_GoTo_Off);
            }
            MI->Open_Buffer_Position_Set(Curl_Data.File_GoTo);
            Curl_Data.File_GoTo=(int64u)-1;
        }

        //Parsing
        Result=curl_easy_perform(Curl_Data.Curl);
    }
    while (Result==CURLE_WRITE_ERROR && Curl_Data.File_GoTo!=(int64u)-1);

    #ifdef MEDIAINFO_DEBUG
        std::cout<<std::hex<<Curl_Data.File_Offset-Curl_Data.Debug_BytesRead<<" - "<<Curl_Data.File_Offset<<" : "<<std::dec<<Curl_Data.Debug_BytesRead<<" bytes"<<std::endl;
        std::cout<<"Total: "<<std::dec<<Curl_Data.Debug_BytesRead_Total<<" bytes in "<<Curl_Data.Debug_Count<<" blocks"<<std::endl;
    #endif //MEDIAINFO_DEBUG

    MI->Open_Buffer_Finalize();

    //Cleanup
    curl_easy_cleanup(Curl_Data.Curl);
    return 1;
}

} //NameSpace

#endif //MEDIAINFO_LIBCURL_YES

<|endoftext|>