<commit_before>/******************************************************************************
* Filename: MFListTransforms.cpp
*
* Description:
* This file contains a C++ console application that attempts to list the Media
* Foundation Transforms (MFT) available on the system.
* See https://docs.microsoft.com/en-us/windows/win32/medfound/registering-and-enumerating-mfts.
*
* Author:
* Aaron Clauson (aaron@sipsorcery.com)
*
* History:
* 26 Feb 2015	  Aaron Clauson	  Created, Hobart, Australia.
* 04 Jan 2020   Aaron Clauson   Tidied up and got ti doing something useful.
*
* License: Public Domain (no warranty, use at own risk)
/******************************************************************************/

#include "../Common/MFUtility.h"

#include <stdio.h>
#include <tchar.h>
#include <evr.h>
#include <mfapi.h>
#include <mfplay.h>
#include <mfreadwrite.h>
#include <wmcodecdsp.h>

#pragma comment(lib, "mf.lib")
#pragma comment(lib, "mfplat.lib")
#pragma comment(lib, "mfplay.lib")
#pragma comment(lib, "mfreadwrite.lib")
#pragma comment(lib, "mfuuid.lib")

HRESULT DisplayMFT(IMFActivate* pMFActivate);

int _tmain(int argc, _TCHAR* argv[])
{
	CHECK_HR(CoInitializeEx(NULL, COINIT_APARTMENTTHREADED | COINIT_DISABLE_OLE1DDE),
		"COM initialisation failed.");

	CHECK_HR(MFStartup(MF_VERSION),
		"Media Foundation initialisation failed.");

	IMFActivate** ppActivate = NULL;
	IMFTransform* pDecoder = NULL;
	UINT32 mftCount = 0;

	MFT_REGISTER_TYPE_INFO info = { 0 };
	info.guidMajorType = MFMediaType_Video;
	info.guidSubtype = MFVideoFormat_YUY2;

	MFT_REGISTER_TYPE_INFO outInfo = { 0 };
	outInfo.guidMajorType = MFMediaType_Video;
	outInfo.guidSubtype = MFVideoFormat_H264;

	CHECK_HR(MFTEnumEx(MFT_CATEGORY_VIDEO_ENCODER,
		MFT_ENUM_FLAG_SYNCMFT | MFT_ENUM_FLAG_ASYNCMFT | MFT_ENUM_FLAG_HARDWARE | MFT_ENUM_FLAG_SORTANDFILTER,
		&info,      // Input type
		&outInfo,   // Output type
		&ppActivate,
		&mftCount
		), "MFTEnumEx failed.");

	printf("MFT count %d.\n", mftCount);

	for (int i = 0; i < mftCount; i++) {

		//CHECK_HR(ppActivate[i]->ActivateObject(IID_PPV_ARGS(&pDecoder)),
		//	"Failed to activate decoder.");
    auto hr = DisplayMFT(ppActivate[i]);
	}

done:

	for (UINT32 i = 0; i < mftCount; i++)
	{
		ppActivate[i]->Release();
	}
	CoTaskMemFree(ppActivate);

	printf("finished.\n");
	auto c = getchar();

	return 0;
}

/**
* Prints out the friendly name for a Media Foundation Transform.
*
* Remarks:
* Copied from https://github.com/mvaneerde/blog/blob/develop/mftenum/mftenum/mftenum.cpp.
*/
HRESULT DisplayMFT(IMFActivate* pMFActivate) {
  HRESULT hr;

  // get the CLSID GUID from the IMFAttributes of the activation object
  GUID guidMFT = { 0 };
  hr = pMFActivate->GetGUID(MFT_TRANSFORM_CLSID_Attribute, &guidMFT);
  if (MF_E_ATTRIBUTENOTFOUND == hr) {
    std::cout << "IMFTransform has no CLSID." << std::endl;
    return hr;
  }
  else if (FAILED(hr)) {
    std::cerr << "IMFAttributes::GetGUID(MFT_TRANSFORM_CLSID_Attribute) failed: hr = " << hr << std::endl;
    return hr;
  }

  LPWSTR szGuid = NULL;
  hr = StringFromIID(guidMFT, &szGuid);
  if (FAILED(hr)) {
    std::cerr << "StringFromIID failed: hr = " << hr << std::endl;
    return hr;
  }

  // get the friendly name string from the IMFAttributes of the activation object
  LPWSTR szFriendlyName = NULL;
#pragma prefast(suppress: __WARNING_PASSING_FUNCTION_UNEXPECTED_NULL, "IMFAttributes::GetAllocatedString third argument is optional");
  hr = pMFActivate->GetAllocatedString(
    MFT_FRIENDLY_NAME_Attribute,
    &szFriendlyName,
    NULL // don't care about the length of the string, and MSDN agrees, but SAL annotation is missing "opt"
  );

  if (MF_E_ATTRIBUTENOTFOUND == hr) {
    std::cout << "IMFTransform has no friendly name." << std::endl;
    return hr;
  }
  else if (FAILED(hr)) {
    std::cerr << "IMFAttributes::GetAllocatedString(MFT_FRIENDLY_NAME_Attribute) failed: hr = " << hr << std::endl;
    return hr;
  }
  //CoTaskMemFreeOnExit freeFriendlyName(szFriendlyName);

  std::wcout << szFriendlyName << " (" << szGuid << ")" << std::endl;

  CoTaskMemFree(szGuid);
  CoTaskMemFree(szFriendlyName);

  return S_OK;
}<commit_msg>Listed audio MFT's.<commit_after>/******************************************************************************
* Filename: MFListTransforms.cpp
*
* Description:
* This file contains a C++ console application that attempts to list the Media
* Foundation Transforms (MFT) available on the system.
* See https://docs.microsoft.com/en-us/windows/win32/medfound/registering-and-enumerating-mfts.
*
* Author:
* Aaron Clauson (aaron@sipsorcery.com)
*
* History:
* 26 Feb 2015	  Aaron Clauson	  Created, Hobart, Australia.
* 04 Jan 2020   Aaron Clauson   Tidied up and got it doing something useful.
*
* License: Public Domain (no warranty, use at own risk)
/******************************************************************************/

#include "../Common/MFUtility.h"

#include <stdio.h>
#include <tchar.h>
#include <evr.h>
#include <mfapi.h>
#include <mfplay.h>
#include <mfreadwrite.h>
#include <wmcodecdsp.h>

#pragma comment(lib, "mf.lib")
#pragma comment(lib, "mfplat.lib")
#pragma comment(lib, "mfplay.lib")
#pragma comment(lib, "mfreadwrite.lib")
#pragma comment(lib, "mfuuid.lib")

void ListTansforms(MFT_REGISTER_TYPE_INFO* pInput, MFT_REGISTER_TYPE_INFO* pOutput);
HRESULT DisplayMFT(IMFActivate* pMFActivate);

int _tmain(int argc, _TCHAR* argv[])
{
	CHECK_HR(CoInitializeEx(NULL, COINIT_APARTMENTTHREADED | COINIT_DISABLE_OLE1DDE),
		"COM initialisation failed.");

	CHECK_HR(MFStartup(MF_VERSION),
		"Media Foundation initialisation failed.");

  // Audio subtype GUIDS : https://docs.microsoft.com/en-us/windows/win32/medfound/audio-subtype-guids
  // Video subtype GUIDS : https://docs.microsoft.com/en-us/windows/win32/medfound/video-subtype-guids

  MFT_REGISTER_TYPE_INFO videoYuv = { MFMediaType_Video, MFVideoFormat_YUY2 };
  MFT_REGISTER_TYPE_INFO videoRgb24 = { MFMediaType_Video, MFVideoFormat_RGB24 };
  MFT_REGISTER_TYPE_INFO videoH264 = { MFMediaType_Video, MFVideoFormat_H264 };
  MFT_REGISTER_TYPE_INFO videoVP8 = { MFMediaType_Video, MFVideoFormat_VP80 };
	MFT_REGISTER_TYPE_INFO audioPcm = { MFMediaType_Audio, MFAudioFormat_PCM };
	MFT_REGISTER_TYPE_INFO audioMp3 = { MFMediaType_Audio, MFAudioFormat_MP3 };

  std::cout << "Video MFT's for input: YUV2, output: all" << std::endl;
  ListTansforms(&videoYuv, NULL);
  std::cout << std::endl;

  std::cout << "Video MFT's for input: YUV2, output: H264" << std::endl;
  ListTansforms(&videoYuv, &videoH264);
  std::cout << std::endl;

  std::cout << "Video MFT's for input: YUV2, output: VP8" << std::endl;
  ListTansforms(&videoYuv, &videoVP8);
  std::cout << std::endl;

  std::cout << "Video MFT's for input: RGB24, output: H264" << std::endl;
  ListTansforms(&videoRgb24, &videoH264);
  std::cout << std::endl;

  std::cout << "Audio MFT's for input: PCM, output: all" << std::endl;
  ListTansforms(&audioPcm, NULL);
  std::cout << std::endl;

  std::cout << "Audio MFT's for input: PCM, output: MP3" << std::endl;
  ListTansforms(&audioPcm, &audioMp3);
  std::cout << std::endl;

done:
	printf("finished.\n");
	auto c = getchar();

	return 0;
}

/**
* Prints out a list of all the Media Foundation Transforms that match for
* the input media type and the optional output media type.
* @param[in] pInput: pointer to the MFT input type to match.
* @param[in] pOutput: Optional. Pointer to the MFT output type to match. If
                      NULL all MFT's for teh input type will be listed.
* Remarks:
* Copied from https://github.com/mvaneerde/blog/blob/develop/mftenum/mftenum/mftenum.cpp.
*/
void ListTansforms(MFT_REGISTER_TYPE_INFO* pInput, MFT_REGISTER_TYPE_INFO* pOutput)
{
  IMFActivate** ppActivate = NULL;
  IMFTransform* pDecoder = NULL;
  UINT32 mftCount = 0;

  auto category = (pInput->guidMajorType == MFMediaType_Audio) ? MFT_CATEGORY_AUDIO_ENCODER : MFT_CATEGORY_VIDEO_ENCODER;
  CHECK_HR(MFTEnumEx(category,
    NULL, //MFT_ENUM_FLAG_SYNCMFT | MFT_ENUM_FLAG_ASYNCMFT | MFT_ENUM_FLAG_HARDWARE | MFT_ENUM_FLAG_SORTANDFILTER,
    pInput,
    pOutput,
    &ppActivate,
    &mftCount
  ), "MFTEnumEx failed.");

  printf("MFT count %d.\n", mftCount);

  for (int i = 0; i < mftCount; i++) {
    auto hr = DisplayMFT(ppActivate[i]);
  }

done:

  for (UINT32 i = 0; i < mftCount; i++)
  {
    ppActivate[i]->Release();
  }
  CoTaskMemFree(ppActivate);
}

/**
* Prints out the friendly name for a Media Foundation Transform.
*
* Remarks:
* Copied from https://github.com/mvaneerde/blog/blob/develop/mftenum/mftenum/mftenum.cpp.
*/
HRESULT DisplayMFT(IMFActivate* pMFActivate) {
  HRESULT hr;

  // get the CLSID GUID from the IMFAttributes of the activation object
  GUID guidMFT = { 0 };
  hr = pMFActivate->GetGUID(MFT_TRANSFORM_CLSID_Attribute, &guidMFT);
  if (MF_E_ATTRIBUTENOTFOUND == hr) {
    std::cout << "IMFTransform has no CLSID." << std::endl;
    return hr;
  }
  else if (FAILED(hr)) {
    std::cerr << "IMFAttributes::GetGUID(MFT_TRANSFORM_CLSID_Attribute) failed: hr = " << hr << std::endl;
    return hr;
  }

  LPWSTR szGuid = NULL;
  hr = StringFromIID(guidMFT, &szGuid);
  if (FAILED(hr)) {
    std::cerr << "StringFromIID failed: hr = " << hr << std::endl;
    return hr;
  }

  // get the friendly name string from the IMFAttributes of the activation object
  LPWSTR szFriendlyName = NULL;
  UINT len = 0;
  hr = pMFActivate->GetAllocatedString(
    MFT_FRIENDLY_NAME_Attribute,
    &szFriendlyName,
    &len
  );

  if (MF_E_ATTRIBUTENOTFOUND == hr) {
    std::cout << "IMFTransform has no friendly name." << std::endl;
    return hr;
  }
  else if (FAILED(hr)) {
    std::cerr << "IMFAttributes::GetAllocatedString(MFT_FRIENDLY_NAME_Attribute) failed: hr = " << hr << std::endl;
    return hr;
  }
  //CoTaskMemFreeOnExit freeFriendlyName(szFriendlyName);

  std::wcout << szFriendlyName << " (" << szGuid << ")" << std::endl;

  CoTaskMemFree(szGuid);
  CoTaskMemFree(szFriendlyName);

  return S_OK;
}<|endoftext|>


<commit_before>/* Process all images in the directory specified in the argument.
Xiang Xiang (eglxiang@gmail.com), May 2016, MIT license.
Main functionality: compute VGG_Face features for each image.
First you need to create the saving directory yourself.
*/

#define CPU_ONLY
#define CORR_METRIC
#include "Classifier.h"
#include <caffe/caffe.hpp>
#include <opencv2/core/core.hpp>
#include <opencv2/highgui/highgui.hpp>
#include <opencv2/imgproc/imgproc.hpp>
#include <iostream>
#include <math.h>
#include <dirent.h>
#include <string.h>
//#include <boost/python.hpp>

using namespace caffe;
using namespace cv;
using namespace std;
//using namespace boost::python

int main(int argc, char** argv)
{
    // load image
    if (argc != 4)
    {
        cerr<<"Run a simple test sample using pretrained VGG face model. " << endl
            << "Usage: " << argv[0]
            << "deploy.prototxt network.caffemodel folderpath" << endl;
        return -1;
    }
    string model_file = argv[1];
    string trained_file = argv[2];
    Classifier classifier(model_file, trained_file);
    string foldername = argv[3];
    string rootpath = "/mnt/localsata/selected_faces/";
    string folderpath = rootpath + foldername;

    // list files in the folder
    DIR *dir;
    dir = opendir(folderpath.c_str());
    cout << dir << endl;
    string imgName;
    struct dirent *ent;

    if (dir != NULL) {
        while ( (ent = readdir(dir)) != NULL ) {
            imgName = ent->d_name;
            if (imgName.compare(".")!=0 && imgName.compare("..")!=0)
            {
                string aux;
                aux.append(folderpath);
                aux.append(imgName);
                cout << aux << endl;
                Mat img = imread(aux, -1);
                if (! img.data) {
                    cout << "Could not open or find the image" << endl;
                    return -1;
                }
                int height = img.rows;
                int width = img.cols;

                // subtract the average face
                Mat avg(height, width, CV_8UC3, Scalar(93.5940,104.7624,129.1863));
                // Mat mean_img(height, width, CV_8UC3, Scalar(90,100,120)); //also works
                Mat image = img - avg;
                vector<float> output = classifier.Predict(image);

                //for (vector<float>::const_iterator i = output.begin(); i != output.end(); ++i) {
                //    cout << *i << ' ';
                //}
            }
        }
    }

    return 0;
}
<commit_msg>Add comment<commit_after>/* Process all images in the directory specified in the argument.
Xiang Xiang (eglxiang@gmail.com), May 2016, MIT license.
Main functionality: compute VGG_Face features for each image.
First you need to create the saving directory yourself.
It writes the feature vector of each face image into a txt file.
*/

#define CPU_ONLY
#define CORR_METRIC
#include "Classifier.h"
#include <caffe/caffe.hpp>
#include <opencv2/core/core.hpp>
#include <opencv2/highgui/highgui.hpp>
#include <opencv2/imgproc/imgproc.hpp>
#include <iostream>
#include <math.h>
#include <dirent.h>
#include <string.h>
//#include <boost/python.hpp>

using namespace caffe;
using namespace cv;
using namespace std;
//using namespace boost::python

int main(int argc, char** argv)
{
    // load image
    if (argc != 4)
    {
        cerr<<"Run a simple test sample using pretrained VGG face model. " << endl
            << "Usage: " << argv[0]
            << "deploy.prototxt network.caffemodel folderpath" << endl;
        return -1;
    }
    string model_file = argv[1];
    string trained_file = argv[2];
    Classifier classifier(model_file, trained_file);
    string foldername = argv[3];
    string rootpath = "/mnt/localsata/selected_faces/";
    string folderpath = rootpath + foldername;

    // list files in the folder
    DIR *dir;
    dir = opendir(folderpath.c_str());
    cout << dir << endl;
    string imgName;
    struct dirent *ent;

    if (dir != NULL) {
        while ( (ent = readdir(dir)) != NULL ) {
            imgName = ent->d_name;
            if (imgName.compare(".")!=0 && imgName.compare("..")!=0)
            {
                string aux;
                aux.append(folderpath);
                aux.append(imgName);
                cout << aux << endl;
                Mat img = imread(aux, -1);
                if (! img.data) {
                    cout << "Could not open or find the image" << endl;
                    return -1;
                }
                int height = img.rows;
                int width = img.cols;

                // subtract the average face
                Mat avg(height, width, CV_8UC3, Scalar(93.5940,104.7624,129.1863));
                // Mat mean_img(height, width, CV_8UC3, Scalar(90,100,120)); //also works
                Mat image = img - avg;
                vector<float> output = classifier.Predict(image);

                //for (vector<float>::const_iterator i = output.begin(); i != output.end(); ++i) {
                //    cout << *i << ' ';
                //}
            }
        }
    }

    return 0;
}
<|endoftext|>


<commit_before>/*
    virtual_keyboard.cpp:

    Copyright (C) 2006 Steven Yi

    This file is part of Csound.

    The Csound 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 2.1 of the License, or (at your option) any later version.

    Csound 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 Csound; if not, write to the Free Software
    Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA
    02111-1307 USA
*/

#include "csdl.h"
#include "csGblMtx.h"
#include "midiops.h"
#include "oload.h"
#include "winFLTK.h"
#include "FLTKKeyboardWindow.hpp"
#include <map>

static FLTKKeyboardWindow *createWindow(CSOUND *csound) {
       return new FLTKKeyboardWindow(csound, 624, 120, "Csound Virtual Keyboard");
}

static void deleteWindow(CSOUND *csound, FLTKKeyboardWindow * keyWin) {
	if(keyWin == NULL) {
		return;
	}
	Fl_lock(csound);

	keyWin->hide();
    delete keyWin;

	Fl_awake(csound);
	Fl_wait(csound, 0.0);
	Fl_unlock(csound);
}

extern "C"
{

static int OpenMidiInDevice_(CSOUND *csound, void **userData, const char *dev)
{
    FLTKKeyboardWindow *keyboard = createWindow(csound);
    *userData = (void *)keyboard;

    Fl_lock(csound);
    keyboard->show();
    Fl_wait(csound, 0.0);
    Fl_unlock(csound);

    /* report success */
    return 0;
}

static int OpenMidiOutDevice_(CSOUND *csound, void **userData, const char *dev)
{
    /* report success */
    return 0;
}

static int ReadMidiData_(CSOUND *csound, void *userData,
                         unsigned char *mbuf, int nbytes)
{

    FLTKKeyboardWindow *keyWin = (FLTKKeyboardWindow *)userData;

    Fl_lock(csound);
	Fl_awake(csound);
	Fl_wait(csound, 0.0);
	Fl_unlock(csound);

    if(!keyWin->visible()) {
        return 0;
    }

	int count = 0;
	int channel = keyWin->channel;

	keyWin->lock();

	if(keyWin->previousProgram[channel] != keyWin->program[channel]) {
		*mbuf++ = (unsigned char)(0xC0 + channel);
        *mbuf++ = (unsigned char)(keyWin->program[channel]);

        count += 2;

        keyWin->previousProgram[channel] = keyWin->program[channel];
	}

	keyWin->unlock();

	keyWin->keyboard->lock();

	int *changedKeyStates = keyWin->keyboard->changedKeyStates;
 	int *keyStates = keyWin->keyboard->keyStates;

    for(int i = 0; i < 88; i++) {
    	if(keyStates[i] == -1) {
    		*mbuf++ = (unsigned char)0x90 + channel;
            *mbuf++ = (unsigned char)i + 21;
            *mbuf++ = (unsigned char)0;

            count += 3;
            keyStates[i] = 0;
    	} else if(changedKeyStates[i] != keyStates[i]) {

	        if(keyStates[i] == 1) {
	            *mbuf++ = (unsigned char)0x90 + channel;
	            *mbuf++ = (unsigned char)i + 21;
	            *mbuf++ = (unsigned char)127;

	            count += 3;

	        } else {
	            *mbuf++ = (unsigned char)0x90 + channel;
	            *mbuf++ = (unsigned char)i + 21;
	            *mbuf++ = (unsigned char)0;

	            count += 3;
	        }
    	}

		changedKeyStates[i] = keyStates[i];
    }

    if(keyWin->keyboard->aNotesOff == 1) {
    	keyWin->keyboard->aNotesOff = 0;
    	*mbuf++ = (unsigned char)0xB0;
        *mbuf++ = (unsigned char)123;
        *mbuf++ = (unsigned char)0;

        count += 3;
    }

	keyWin->keyboard->unlock();

    return count;
}

static int WriteMidiData_(CSOUND *csound, void *userData,
                          const unsigned char *mbuf, int nbytes)
{
    /* return the number of bytes written */
    return 0;
}

static int CloseMidiInDevice_(CSOUND *csound, void *userData)
{

    deleteWindow(csound, (FLTKKeyboardWindow *)userData);

    return 0;
}

static int CloseMidiOutDevice_(CSOUND *csound, void *userData)
{
    return 0;
}

/* module interface functions */

PUBLIC int csoundModuleCreate(CSOUND *csound)
{
    /* nothing to do, report success */
    csound->Message(csound, "virtual_keyboard real time MIDI plugin for Csound\n");
    return 0;
}

PUBLIC int csoundModuleInit(CSOUND *csound)
{
    char    *drv;

    if (csound->QueryGlobalVariable(csound,
                                    "FLTK_Flags") == (void*) 0) {
      if (csound->CreateGlobalVariable(csound,
                                       "FLTK_Flags", sizeof(int)) != 0)
        csound->Die(csound,
                    Str("virtual_keyboard.cpp: error allocating FLTK flags"));
    }

    drv = (char*) (csound->QueryGlobalVariable(csound, "_RTMIDI"));
    if (drv == NULL)
      return 0;
    if (!(strcmp(drv, "virtual") == 0))
      return 0;
    csound->Message(csound, "rtmidi: virtual_keyboard module enabled\n");
    csound->SetExternalMidiInOpenCallback(csound, OpenMidiInDevice_);
    csound->SetExternalMidiReadCallback(csound, ReadMidiData_);
    csound->SetExternalMidiInCloseCallback(csound, CloseMidiInDevice_);
    csound->SetExternalMidiOutOpenCallback(csound, OpenMidiOutDevice_);
    csound->SetExternalMidiWriteCallback(csound, WriteMidiData_);
    csound->SetExternalMidiOutCloseCallback(csound, CloseMidiOutDevice_);

    return 0;
}

PUBLIC int csoundModuleInfo(void)
{
    /* does not depend on MYFLT type */
    return ((CS_APIVERSION << 16) + (CS_APISUBVER << 8));
}

} // END EXTERN C
<commit_msg>removed unnecessary header<commit_after>/*
    virtual_keyboard.cpp:

    Copyright (C) 2006 Steven Yi

    This file is part of Csound.

    The Csound 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 2.1 of the License, or (at your option) any later version.

    Csound 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 Csound; if not, write to the Free Software
    Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA
    02111-1307 USA
*/

#include "csdl.h"
#include "csGblMtx.h"
#include "midiops.h"
#include "oload.h"
#include "winFLTK.h"
#include "FLTKKeyboardWindow.hpp"

static FLTKKeyboardWindow *createWindow(CSOUND *csound) {
       return new FLTKKeyboardWindow(csound, 624, 120, "Csound Virtual Keyboard");
}

static void deleteWindow(CSOUND *csound, FLTKKeyboardWindow * keyWin) {
	if(keyWin == NULL) {
		return;
	}
	Fl_lock(csound);

	keyWin->hide();
    delete keyWin;

	Fl_awake(csound);
	Fl_wait(csound, 0.0);
	Fl_unlock(csound);
}

extern "C"
{

static int OpenMidiInDevice_(CSOUND *csound, void **userData, const char *dev)
{
    FLTKKeyboardWindow *keyboard = createWindow(csound);
    *userData = (void *)keyboard;

    Fl_lock(csound);
    keyboard->show();
    Fl_wait(csound, 0.0);
    Fl_unlock(csound);

    /* report success */
    return 0;
}

static int OpenMidiOutDevice_(CSOUND *csound, void **userData, const char *dev)
{
    /* report success */
    return 0;
}

static int ReadMidiData_(CSOUND *csound, void *userData,
                         unsigned char *mbuf, int nbytes)
{

    FLTKKeyboardWindow *keyWin = (FLTKKeyboardWindow *)userData;

    Fl_lock(csound);
	Fl_awake(csound);
	Fl_wait(csound, 0.0);
	Fl_unlock(csound);

    if(!keyWin->visible()) {
        return 0;
    }

	int count = 0;
	int channel = keyWin->channel;

	keyWin->lock();

	if(keyWin->previousProgram[channel] != keyWin->program[channel]) {
		*mbuf++ = (unsigned char)(0xC0 + channel);
        *mbuf++ = (unsigned char)(keyWin->program[channel]);

        count += 2;

        keyWin->previousProgram[channel] = keyWin->program[channel];
	}

	keyWin->unlock();

	keyWin->keyboard->lock();

	int *changedKeyStates = keyWin->keyboard->changedKeyStates;
 	int *keyStates = keyWin->keyboard->keyStates;

    for(int i = 0; i < 88; i++) {
    	if(keyStates[i] == -1) {
    		*mbuf++ = (unsigned char)0x90 + channel;
            *mbuf++ = (unsigned char)i + 21;
            *mbuf++ = (unsigned char)0;

            count += 3;
            keyStates[i] = 0;
    	} else if(changedKeyStates[i] != keyStates[i]) {

	        if(keyStates[i] == 1) {
	            *mbuf++ = (unsigned char)0x90 + channel;
	            *mbuf++ = (unsigned char)i + 21;
	            *mbuf++ = (unsigned char)127;

	            count += 3;

	        } else {
	            *mbuf++ = (unsigned char)0x90 + channel;
	            *mbuf++ = (unsigned char)i + 21;
	            *mbuf++ = (unsigned char)0;

	            count += 3;
	        }
    	}

		changedKeyStates[i] = keyStates[i];
    }

    if(keyWin->keyboard->aNotesOff == 1) {
    	keyWin->keyboard->aNotesOff = 0;
    	*mbuf++ = (unsigned char)0xB0;
        *mbuf++ = (unsigned char)123;
        *mbuf++ = (unsigned char)0;

        count += 3;
    }

	keyWin->keyboard->unlock();

    return count;
}

static int WriteMidiData_(CSOUND *csound, void *userData,
                          const unsigned char *mbuf, int nbytes)
{
    /* return the number of bytes written */
    return 0;
}

static int CloseMidiInDevice_(CSOUND *csound, void *userData)
{

    deleteWindow(csound, (FLTKKeyboardWindow *)userData);

    return 0;
}

static int CloseMidiOutDevice_(CSOUND *csound, void *userData)
{
    return 0;
}

/* module interface functions */

PUBLIC int csoundModuleCreate(CSOUND *csound)
{
    /* nothing to do, report success */
    csound->Message(csound, "virtual_keyboard real time MIDI plugin for Csound\n");
    return 0;
}

PUBLIC int csoundModuleInit(CSOUND *csound)
{
    char    *drv;

    if (csound->QueryGlobalVariable(csound,
                                    "FLTK_Flags") == (void*) 0) {
      if (csound->CreateGlobalVariable(csound,
                                       "FLTK_Flags", sizeof(int)) != 0)
        csound->Die(csound,
                    Str("virtual_keyboard.cpp: error allocating FLTK flags"));
    }

    drv = (char*) (csound->QueryGlobalVariable(csound, "_RTMIDI"));
    if (drv == NULL)
      return 0;
    if (!(strcmp(drv, "virtual") == 0))
      return 0;
    csound->Message(csound, "rtmidi: virtual_keyboard module enabled\n");
    csound->SetExternalMidiInOpenCallback(csound, OpenMidiInDevice_);
    csound->SetExternalMidiReadCallback(csound, ReadMidiData_);
    csound->SetExternalMidiInCloseCallback(csound, CloseMidiInDevice_);
    csound->SetExternalMidiOutOpenCallback(csound, OpenMidiOutDevice_);
    csound->SetExternalMidiWriteCallback(csound, WriteMidiData_);
    csound->SetExternalMidiOutCloseCallback(csound, CloseMidiOutDevice_);

    return 0;
}

PUBLIC int csoundModuleInfo(void)
{
    /* does not depend on MYFLT type */
    return ((CS_APIVERSION << 16) + (CS_APISUBVER << 8));
}

} // END EXTERN C
<|endoftext|>


<commit_before>/*=========================================================================

Program:   Medical Imaging & Interaction Toolkit
Language:  C++
Date:      $Date: 2009-05-05 19:02:32 +0200 (Tue, 05 May 2009) $
Version:   $Revision: 17105 $

Copyright (c) German Cancer Research Center, Division of Medical and
Biological Informatics. All rights reserved.
See MITKCopyright.txt or http://www.mitk.org/copyright.html for details.

This software is distributed WITHOUT ANY WARRANTY; without even
the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR
PURPOSE.  See the above copyright notices for more information.

=========================================================================*/

#include "itkImageFileWriter.h"
#include "itkWarpImageFilter.h"
#include "itkImageRegionIterator.h"

#include "mitkBSplineRegistration.h"
#include "itkBSplineDeformableTransform.h"
//#include "itkLBFGSOptimizer.h"
#include "itkMeanSquaresImageToImageMetric.h"
#include "itkMattesMutualInformationImageToImageMetric.h"
#include "itkResampleImageFilter.h"
#include "itkImageRegistrationMethod.h"

#include "itkBSplineDeformableTransformInitializer.h"

#include "mitkOptimizerFactory.h"

namespace mitk {

  BSplineRegistration::BSplineRegistration():
    m_Iterations(50),
    m_ResultName("deformedImage.mhd"),
    m_SaveResult(true),
    m_SaveDeformationField(false),
    m_UpdateInputImage(false)
  {

  }

  BSplineRegistration::~BSplineRegistration()
  {
  }

  void BSplineRegistration::SetNumberOfIterations(int iterations)
  {
    m_Iterations = iterations;
  }

  void BSplineRegistration::SetSaveResult(bool saveResult)
  {
    m_SaveResult = saveResult;
  }

  void BSplineRegistration::SetResultFileName(const char* resultName)
  {
    m_ResultName = resultName;
  }

  template < typename TPixel, unsigned int VImageDimension >
    void BSplineRegistration::GenerateData2( itk::Image<TPixel, VImageDimension>* itkImage1)
  {
    std::cout << "start bspline registration" << std::endl;

    // Typedefs
    typedef typename itk::Image< TPixel, VImageDimension >  FixedImageType;
    typedef typename itk::Image< TPixel, VImageDimension >  MovingImageType;

    typedef typename itk::Vector< float, VImageDimension >    VectorPixelType;
    typedef typename itk::Image<  VectorPixelType, VImageDimension > DeformationFieldType;

    typedef itk::BSplineDeformableTransform<
                                double,
                                VImageDimension,
                                3 >                         TransformType;

    typedef typename TransformType::ParametersType          ParametersType;

    //typedef itk::LBFGSOptimizer                             OptimizerType;
    typedef itk::SingleValuedNonLinearOptimizer             OptimizerType;
    typedef itk::MattesMutualInformationImageToImageMetric<
                                FixedImageType,
                                MovingImageType >           MetricType;

    /*typedef itk::MeanSquaresImageToImageMetric<
                                FixedImageType,
                                MovingImageType >           MetricType;*/

    typedef itk::LinearInterpolateImageFunction<
                                MovingImageType,
                                double >                    InterpolatorType;

    typedef itk::ImageRegistrationMethod<
                                FixedImageType,
                                MovingImageType >           RegistrationType;

    typedef typename itk::WarpImageFilter<
                            MovingImageType,
                            MovingImageType,
                            DeformationFieldType  >         WarperType;

    typedef typename TransformType::SpacingType                      SpacingType;

    typedef typename TransformType::OriginType                       OriginType;

    typedef itk::ResampleImageFilter<
                                MovingImageType,
                                FixedImageType >            ResampleFilterType;

    typedef itk::Image< TPixel, VImageDimension >           OutputImageType;

    // Sample new image with the same image type as the fixed image
    typedef itk::CastImageFilter<
                                FixedImageType,
                                FixedImageType >            CastFilterType;

    typedef itk::Vector< float, VImageDimension >           VectorType;
    typedef itk::Image< VectorType, VImageDimension >       DeformationFieldType;

    typedef itk::BSplineDeformableTransformInitializer <
                                TransformType,
                                FixedImageType >            InitializerType;

    typename InterpolatorType::Pointer   interpolator  = InterpolatorType::New();
    typename RegistrationType::Pointer   registration  = RegistrationType::New();
    typename InitializerType::Pointer    initializer   = InitializerType::New();
    typename TransformType::Pointer      transform     = TransformType::New();
    typename MetricType::Pointer         metric        = MetricType::New();

    metric->SetNumberOfHistogramBins( 24);
    metric->SetNumberOfSpatialSamples(100000);

    typename OptimizerFactory::Pointer optFac = OptimizerFactory::New();
    optFac->SetOptimizerParameters(m_OptimizerParameters);
    optFac->SetNumberOfTransformParameters(transform->GetNumberOfParameters());
    OptimizerType::Pointer optimizer = optFac->GetOptimizer();

    typename FixedImageType::Pointer fixedImage = FixedImageType::New();
    mitk::CastToItkImage(m_ReferenceImage, fixedImage);
    typename MovingImageType::Pointer movingImage = itkImage1;
    typename FixedImageType::RegionType fixedRegion = fixedImage->GetBufferedRegion();
    typename MovingImageType::RegionType movingRegion = movingImage->GetBufferedRegion();

    registration->SetMetric(          metric        );
    registration->SetOptimizer(       optimizer     );
    registration->SetInterpolator(    interpolator  );
    registration->SetFixedImage(      fixedImage    );
    registration->SetMovingImage(     movingImage   );
    registration->SetFixedImageRegion(fixedRegion   );

    initializer->SetTransform(transform);
    initializer->SetImage(fixedImage);
    initializer->SetNumberOfGridNodesInsideTheImage( m_NumberOfGridPoints );
    initializer->InitializeTransform();

    registration->SetTransform( transform );

    const unsigned int numberOfParameters = transform->GetNumberOfParameters();

    typename itk::BSplineDeformableTransform<
                                double,
                                VImageDimension,
                                3 >::ParametersType  parameters;

    parameters.set_size(numberOfParameters);
    parameters.Fill( 0.0 );
    transform->SetParameters( parameters );

    // We now pass the parameters of the current transform as the initial
    // parameters to be used when the registration process starts.
    registration->SetInitialTransformParameters( transform->GetParameters() );

    std::cout << "Intial Parameters = " << std::endl;
    std::cout << transform->GetParameters() << std::endl;

    std::cout << std::endl << "Starting Registration" << std::endl;

    try
    {
      registration->StartRegistration();
    }
    catch( itk::ExceptionObject & err )
    {
      std::cerr << "ExceptionObject caught !" << std::endl;
      std::cerr << err << std::endl;
    }

    typename OptimizerType::ParametersType finalParameters =
                      registration->GetLastTransformParameters();

    std::cout << "Last Transform Parameters" << std::endl;
    std::cout << finalParameters << std::endl;

    transform->SetParameters( finalParameters );

/*
    ResampleFilterType::Pointer       resampler = ResampleFilterType::New();
    resampler->SetTransform(          transform );
    resampler->SetInput(              movingImage );
    resampler->SetSize(               fixedImage->GetLargestPossibleRegion().GetSize() );
    resampler->SetOutputOrigin(       fixedImage->GetOrigin() );
    resampler->SetOutputSpacing(      fixedImage->GetSpacing() );
    resampler->SetOutputDirection(    fixedImage->GetDirection() );
    resampler->SetDefaultPixelValue(  100 );
    resampler->SetInterpolator(       interpolator);
    resampler->Update();*/

    // Generate deformation field
    typename DeformationFieldType::Pointer field = DeformationFieldType::New();
    field->SetRegions( movingRegion );
    field->SetOrigin( movingImage->GetOrigin() );
    field->SetSpacing( movingImage->GetSpacing() );
    field->SetDirection( movingImage->GetDirection() );
    field->Allocate();

    typedef itk::ImageRegionIterator< DeformationFieldType > FieldIterator;
    FieldIterator fi( field, movingRegion );
    fi.GoToBegin();

    typename TransformType::InputPointType  fixedPoint;
    typename TransformType::OutputPointType movingPoint;
    typename DeformationFieldType::IndexType index;

    VectorType displacement;

    while( ! fi.IsAtEnd() )
    {
      index = fi.GetIndex();
      field->TransformIndexToPhysicalPoint( index, fixedPoint );
      movingPoint = transform->TransformPoint( fixedPoint );
      displacement = movingPoint - fixedPoint;
      fi.Set( displacement );
      ++fi;
    }

    // Use the deformation field to warp the moving image
    typename WarperType::Pointer warper = WarperType::New();
    warper->SetInput( movingImage );
    warper->SetInterpolator( interpolator );
    warper->SetOutputSpacing( movingImage->GetSpacing() );
    warper->SetOutputOrigin( movingImage->GetOrigin() );
    warper->SetOutputDirection( movingImage->GetDirection() );
    warper->SetDeformationField( field );
    warper->Update();

    typename MovingImageType::Pointer result = warper->GetOutput();

    if(m_UpdateInputImage)
    {
      Image::Pointer outputImage = this->GetOutput();
      mitk::CastToMitkImage( result, outputImage );
    }

    // Save the deformationfield resulting from the registration
    if(m_SaveDeformationField)
    {
      typedef itk::ImageFileWriter< DeformationFieldType >  FieldWriterType;
      typename FieldWriterType::Pointer fieldWriter = FieldWriterType::New();

      fieldWriter->SetInput( field );

      fieldWriter->SetFileName( m_DeformationFileName );
      try
      {
        fieldWriter->Update();
      }
      catch( itk::ExceptionObject & excp )
      {
        std::cerr << "Exception thrown " << std::endl;
        std::cerr << excp << std::endl;
        //return EXIT_FAILURE;
      }
    }

  }
} // end namespace
<commit_msg>CHG (#1927): Changed default number spatial samples in metric<commit_after>/*=========================================================================

Program:   Medical Imaging & Interaction Toolkit
Language:  C++
Date:      $Date: 2009-05-05 19:02:32 +0200 (Tue, 05 May 2009) $
Version:   $Revision: 17105 $

Copyright (c) German Cancer Research Center, Division of Medical and
Biological Informatics. All rights reserved.
See MITKCopyright.txt or http://www.mitk.org/copyright.html for details.

This software is distributed WITHOUT ANY WARRANTY; without even
the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR
PURPOSE.  See the above copyright notices for more information.

=========================================================================*/

#include "itkImageFileWriter.h"
#include "itkWarpImageFilter.h"
#include "itkImageRegionIterator.h"

#include "mitkBSplineRegistration.h"
#include "itkBSplineDeformableTransform.h"
//#include "itkLBFGSOptimizer.h"
#include "itkMeanSquaresImageToImageMetric.h"
#include "itkMattesMutualInformationImageToImageMetric.h"
#include "itkResampleImageFilter.h"
#include "itkImageRegistrationMethod.h"

#include "itkBSplineDeformableTransformInitializer.h"

#include "mitkOptimizerFactory.h"

namespace mitk {

  BSplineRegistration::BSplineRegistration():
    m_Iterations(50),
    m_ResultName("deformedImage.mhd"),
    m_SaveResult(true),
    m_SaveDeformationField(false),
    m_UpdateInputImage(false)
  {

  }

  BSplineRegistration::~BSplineRegistration()
  {
  }

  void BSplineRegistration::SetNumberOfIterations(int iterations)
  {
    m_Iterations = iterations;
  }

  void BSplineRegistration::SetSaveResult(bool saveResult)
  {
    m_SaveResult = saveResult;
  }

  void BSplineRegistration::SetResultFileName(const char* resultName)
  {
    m_ResultName = resultName;
  }

  template < typename TPixel, unsigned int VImageDimension >
    void BSplineRegistration::GenerateData2( itk::Image<TPixel, VImageDimension>* itkImage1)
  {
    std::cout << "start bspline registration" << std::endl;

    // Typedefs
    typedef typename itk::Image< TPixel, VImageDimension >  FixedImageType;
    typedef typename itk::Image< TPixel, VImageDimension >  MovingImageType;

    typedef typename itk::Vector< float, VImageDimension >    VectorPixelType;
    typedef typename itk::Image<  VectorPixelType, VImageDimension > DeformationFieldType;

    typedef itk::BSplineDeformableTransform<
                                double,
                                VImageDimension,
                                3 >                         TransformType;

    typedef typename TransformType::ParametersType          ParametersType;

    //typedef itk::LBFGSOptimizer                             OptimizerType;
    typedef itk::SingleValuedNonLinearOptimizer             OptimizerType;
    typedef itk::MattesMutualInformationImageToImageMetric<
                                FixedImageType,
                                MovingImageType >           MetricType;

    /*typedef itk::MeanSquaresImageToImageMetric<
                                FixedImageType,
                                MovingImageType >           MetricType;*/

    typedef itk::LinearInterpolateImageFunction<
                                MovingImageType,
                                double >                    InterpolatorType;

    typedef itk::ImageRegistrationMethod<
                                FixedImageType,
                                MovingImageType >           RegistrationType;

    typedef typename itk::WarpImageFilter<
                            MovingImageType,
                            MovingImageType,
                            DeformationFieldType  >         WarperType;

    typedef typename TransformType::SpacingType                      SpacingType;

    typedef typename TransformType::OriginType                       OriginType;

    typedef itk::ResampleImageFilter<
                                MovingImageType,
                                FixedImageType >            ResampleFilterType;

    typedef itk::Image< TPixel, VImageDimension >           OutputImageType;

    // Sample new image with the same image type as the fixed image
    typedef itk::CastImageFilter<
                                FixedImageType,
                                FixedImageType >            CastFilterType;

    typedef itk::Vector< float, VImageDimension >           VectorType;
    typedef itk::Image< VectorType, VImageDimension >       DeformationFieldType;

    typedef itk::BSplineDeformableTransformInitializer <
                                TransformType,
                                FixedImageType >            InitializerType;

    typename InterpolatorType::Pointer   interpolator  = InterpolatorType::New();
    typename RegistrationType::Pointer   registration  = RegistrationType::New();
    typename InitializerType::Pointer    initializer   = InitializerType::New();
    typename TransformType::Pointer      transform     = TransformType::New();
    typename MetricType::Pointer         metric        = MetricType::New();

    metric->SetNumberOfHistogramBins( 24);
    metric->SetNumberOfSpatialSamples(10000);

    typename OptimizerFactory::Pointer optFac = OptimizerFactory::New();
    optFac->SetOptimizerParameters(m_OptimizerParameters);
    optFac->SetNumberOfTransformParameters(transform->GetNumberOfParameters());
    OptimizerType::Pointer optimizer = optFac->GetOptimizer();

    typename FixedImageType::Pointer fixedImage = FixedImageType::New();
    mitk::CastToItkImage(m_ReferenceImage, fixedImage);
    typename MovingImageType::Pointer movingImage = itkImage1;
    typename FixedImageType::RegionType fixedRegion = fixedImage->GetBufferedRegion();
    typename MovingImageType::RegionType movingRegion = movingImage->GetBufferedRegion();

    registration->SetMetric(          metric        );
    registration->SetOptimizer(       optimizer     );
    registration->SetInterpolator(    interpolator  );
    registration->SetFixedImage(      fixedImage    );
    registration->SetMovingImage(     movingImage   );
    registration->SetFixedImageRegion(fixedRegion   );

    initializer->SetTransform(transform);
    initializer->SetImage(fixedImage);
    initializer->SetNumberOfGridNodesInsideTheImage( m_NumberOfGridPoints );
    initializer->InitializeTransform();

    registration->SetTransform( transform );

    const unsigned int numberOfParameters = transform->GetNumberOfParameters();

    typename itk::BSplineDeformableTransform<
                                double,
                                VImageDimension,
                                3 >::ParametersType  parameters;

    parameters.set_size(numberOfParameters);
    parameters.Fill( 0.0 );
    transform->SetParameters( parameters );

    // We now pass the parameters of the current transform as the initial
    // parameters to be used when the registration process starts.
    registration->SetInitialTransformParameters( transform->GetParameters() );

    std::cout << "Intial Parameters = " << std::endl;
    std::cout << transform->GetParameters() << std::endl;

    std::cout << std::endl << "Starting Registration" << std::endl;

    try
    {
      double tstart(clock());
      registration->StartRegistration();
      double time = clock() - tstart;
      time = time / CLOCKS_PER_SEC;
      LOG_INFO << "Registration time: " << time;
    }
    catch( itk::ExceptionObject & err )
    {
      std::cerr << "ExceptionObject caught !" << std::endl;
      std::cerr << err << std::endl;
    }

    typename OptimizerType::ParametersType finalParameters =
                      registration->GetLastTransformParameters();

    std::cout << "Last Transform Parameters" << std::endl;
    std::cout << finalParameters << std::endl;

    transform->SetParameters( finalParameters );

/*
    ResampleFilterType::Pointer       resampler = ResampleFilterType::New();
    resampler->SetTransform(          transform );
    resampler->SetInput(              movingImage );
    resampler->SetSize(               fixedImage->GetLargestPossibleRegion().GetSize() );
    resampler->SetOutputOrigin(       fixedImage->GetOrigin() );
    resampler->SetOutputSpacing(      fixedImage->GetSpacing() );
    resampler->SetOutputDirection(    fixedImage->GetDirection() );
    resampler->SetDefaultPixelValue(  100 );
    resampler->SetInterpolator(       interpolator);
    resampler->Update();*/

    // Generate deformation field
    typename DeformationFieldType::Pointer field = DeformationFieldType::New();
    field->SetRegions( movingRegion );
    field->SetOrigin( movingImage->GetOrigin() );
    field->SetSpacing( movingImage->GetSpacing() );
    field->SetDirection( movingImage->GetDirection() );
    field->Allocate();

    typedef itk::ImageRegionIterator< DeformationFieldType > FieldIterator;
    FieldIterator fi( field, movingRegion );
    fi.GoToBegin();

    typename TransformType::InputPointType  fixedPoint;
    typename TransformType::OutputPointType movingPoint;
    typename DeformationFieldType::IndexType index;

    VectorType displacement;

    while( ! fi.IsAtEnd() )
    {
      index = fi.GetIndex();
      field->TransformIndexToPhysicalPoint( index, fixedPoint );
      movingPoint = transform->TransformPoint( fixedPoint );
      displacement = movingPoint - fixedPoint;
      fi.Set( displacement );
      ++fi;
    }

    // Use the deformation field to warp the moving image
    typename WarperType::Pointer warper = WarperType::New();
    warper->SetInput( movingImage );
    warper->SetInterpolator( interpolator );
    warper->SetOutputSpacing( movingImage->GetSpacing() );
    warper->SetOutputOrigin( movingImage->GetOrigin() );
    warper->SetOutputDirection( movingImage->GetDirection() );
    warper->SetDeformationField( field );
    warper->Update();

    typename MovingImageType::Pointer result = warper->GetOutput();

    if(m_UpdateInputImage)
    {
      Image::Pointer outputImage = this->GetOutput();
      mitk::CastToMitkImage( result, outputImage );
    }

    // Save the deformationfield resulting from the registration
    if(m_SaveDeformationField)
    {
      typedef itk::ImageFileWriter< DeformationFieldType >  FieldWriterType;
      typename FieldWriterType::Pointer fieldWriter = FieldWriterType::New();

      fieldWriter->SetInput( field );

      fieldWriter->SetFileName( m_DeformationFileName );
      try
      {
        fieldWriter->Update();
      }
      catch( itk::ExceptionObject & excp )
      {
        std::cerr << "Exception thrown " << std::endl;
        std::cerr << excp << std::endl;
        //return EXIT_FAILURE;
      }
    }

  }
} // end namespace
<|endoftext|>


<commit_before>#include "CPUBipartiteGraphAnnealer.h"
#include <cmath>
#include <float.h>
#include <algorithm>
#include <exception>
#include "CPUFormulas.h"

using namespace sqaod;

template<class real>
CPUBipartiteGraphAnnealer<real>::CPUBipartiteGraphAnnealer() {
    m_ = -1;
    annState_ = annNone;
    annealMethod_ = &CPUBipartiteGraphAnnealer::annealOneStepColoring;
#ifdef _OPENMP
    nProcs_ = omp_get_num_procs();
    log("# processors: %d", nProcs_);
#else
    nProcs_ = 1;
#endif
    random_ = new Random[nProcs_];
}

template<class real>
CPUBipartiteGraphAnnealer<real>::~CPUBipartiteGraphAnnealer() {
    delete [] random_;
}

template<class real>
void CPUBipartiteGraphAnnealer<real>::seed(unsigned long seed) {
    for (int idx = 0; idx < nProcs_; ++idx)
        random_[idx].seed(seed + 17 * idx);
    annState_ |= annRandSeedGiven;
}

template<class real>
void CPUBipartiteGraphAnnealer<real>::selectAlgorithm(Algorithm algo) {
    switch (algo) {
    case algoNaive:
        annealMethod_ = &CPUBipartiteGraphAnnealer::annealOneStepNaive;
        break;
    case algoColoring:
    case algoDefault:
        annealMethod_ = &CPUBipartiteGraphAnnealer::annealOneStepColoring;
        break;
    default:
        log("Uknown algo, %s, defaulting to %s.", algoToName(algo), algoToName(algoColoring));
        annealMethod_ = &CPUBipartiteGraphAnnealer::annealOneStepColoring;
    }
}

template<class real>
Algorithm CPUBipartiteGraphAnnealer<real>::algorithm() const {
    if (annealMethod_ == &CPUBipartiteGraphAnnealer::annealOneStepNaive)
        return algoNaive;
    if (annealMethod_ == &CPUBipartiteGraphAnnealer::annealOneStepColoring)
        return algoColoring;
    abort_("Must not reach here.");
    return algoDefault; /* to suppress warning. */
}

template<class real>
void CPUBipartiteGraphAnnealer<real>::getProblemSize(SizeType *N0, SizeType *N1) const {
    *N0 = N0_;
    *N1 = N1_;
}

template<class real>
void CPUBipartiteGraphAnnealer<real>::setProblem(const Vector &b0, const Vector &b1,
                                                 const Matrix &W, OptimizeMethod om) {
    N0_ = (int)b0.size;
    N1_ = (int)b1.size;
    h0_.resize(N0_);
    h1_.resize(N1_);
    J_.resize(N1_, N0_);
    Vector h0(mapFrom(h0_)), h1(mapFrom(h1_));
    Matrix J(mapFrom(J_));
    BGFuncs<real>::calculate_hJc(&h0, &h1, &J, &c_, b0, b1, W);

    om_ = om;
    if (om_ == optMaximize) {
        h0_ *= real(-1.);
        h1_ *= real(-1.);
        J_ *= real(-1.);
        c_ *= real(-1.);
    }
}

template<class real>
void CPUBipartiteGraphAnnealer<real>::setNumTrotters(SizeType m) {
    throwErrorIf(m <= 0, "# trotters must be a positive integer.");
    m_ = m;
    matQ0_.resize(m_, N0_);
    matQ1_.resize(m_, N1_);
    E_.resize(m_);
    annState_ |= annNTrottersGiven;
}

template<class real>
const BitsPairArray &CPUBipartiteGraphAnnealer<real>::get_x() const {
    return bitsPairX_;
}

template<class real>
void CPUBipartiteGraphAnnealer<real>::set_x(const Bits &x0, const Bits &x1) {
    EigenRowVector ex0 = mapToRowVector(x0).cast<real>();
    EigenRowVector ex1 = mapToRowVector(x1).cast<real>();
    matQ0_ = (ex0.array() * 2 - 1).matrix();
    matQ1_ = (ex1.array() * 2 - 1).matrix();
    annState_ |= annQSet;
}

template<class real>
const VectorType<real> &CPUBipartiteGraphAnnealer<real>::get_E() const {
    return E_;
}

template<class real>
void CPUBipartiteGraphAnnealer<real>::get_hJc(Vector *h0, Vector *h1,
                                              Matrix *J, real *c) const {
    mapToRowVector(*h0) = h0_;
    mapToRowVector(*h1) = h1_;
    mapTo(*J) = J_;
    *c = c_;
}

template<class real>
const BitsPairArray &CPUBipartiteGraphAnnealer<real>::get_q() const {
    return bitsPairQ_;
}

template<class real>
void CPUBipartiteGraphAnnealer<real>::randomize_q() {
    real *q = matQ0_.data();
#ifndef _OPENMP
    {
        Random &random = random_[0];
#else
#pragma omp parallel shared(q)
    {
        sqaod::Random &random = random_[omp_get_thread_num()];
#pragma omp for
#endif
        for (int idx = 0; idx < IdxType(N0_ * m_); ++idx) {
            q[idx] = random.randInt(2) ? real(1.) : real(-1.);
        }
        q = matQ1_.data();
#ifdef _OPENMP
#pragma omp for
#endif
        for (int idx = 0; idx < IdxType(N1_ * m_); ++idx) {
            q[idx] = random.randInt(2) ? real(1.) : real(-1.);
        }
    }
    annState_ |= annQSet;
}

template<class real>
void CPUBipartiteGraphAnnealer<real>::calculate_E() {
    BGFuncs<real>::calculate_E(&E_, mapFrom(h0_), mapFrom(h1_), mapFrom(J_), c_,
                               mapFrom(matQ0_), mapFrom(matQ1_));
    if (om_ == optMaximize)
        mapToRowVector(E_) *= real(-1.);
}

template<class real>
void CPUBipartiteGraphAnnealer<real>::initAnneal() {
    if (!(annState_ & annRandSeedGiven))
        for (int idx = 0; idx < nProcs_; ++idx)
            random_[idx].seed();
    annState_ |= annRandSeedGiven;
    if (!(annState_ & annNTrottersGiven))
        setNumTrotters((N0_ + N1_) / 4);
    annState_ |= annNTrottersGiven;
    if (!(annState_ & annQSet))
        randomize_q();
    annState_ |= annQSet;
}

template<class real>
void CPUBipartiteGraphAnnealer<real>::finAnneal() {
    syncBits();
    calculate_E();
}

template<class real>
void CPUBipartiteGraphAnnealer<real>::annealOneStepNaive(real G, real kT) {
    real twoDivM = real(2.) / real(m_);
    real coef = std::log(std::tanh(G / kT / m_)) / kT;
    Random &random = random_[0];
    int N = N0_ + N1_;
    for (int loop = 0; loop < IdxType(N * m_); ++loop) {
        int x = random.randInt(N);
        int y = random.randInt(m_);
        if (x < N0_) {
            real qyx = matQ0_(y, x);
            real sum = J_.transpose().row(x).dot(matQ1_.row(y));
            real dE = - twoDivM * qyx * (h0_(x) + sum);
            int neibour0 = (y == 0) ? m_ - 1 : y - 1;
            int neibour1 = (y == m_ - 1) ? 0 : y + 1;
            dE -= qyx * (matQ0_(neibour0, x) + matQ0_(neibour1, x)) * coef;
            real threshold = (dE < real(0.)) ? real(1.) : std::exp(-dE / kT);
            if (threshold > random.random<real>())
                matQ0_(y, x) = - qyx;
        }
        else {
            x -= N0_;
            real qyx = matQ1_(y, x);
            real sum = J_.row(x).dot(matQ0_.row(y));
            real dE = - twoDivM * qyx * (h1_(x) + sum);
            int neibour0 = (y == 0) ? m_ - 1 : y - 1;
            int neibour1 = (y == m_ - 1) ? 0 : y + 1;
            dE -= qyx * (matQ1_(neibour0, x) + matQ1_(neibour1, x)) * coef;
            real threshold = (dE < real(0.)) ? real(1.) : std::exp(-dE / kT);
            if (threshold > random.random<real>())
                matQ1_(y, x) = - qyx;
        }
    }
}

template<class real>
void CPUBipartiteGraphAnnealer<real>::annealOneStepColoring(real G, real kT) {
    annealHalfStepColoring(N1_, matQ1_, h1_, J_, matQ0_, G, kT);
    annealHalfStepColoring(N0_, matQ0_, h0_, J_.transpose(), matQ1_, G, kT);
}

template<class real>
void CPUBipartiteGraphAnnealer<real>::
annealHalfStepColoring(int N, EigenMatrix &qAnneal,
                       const EigenRowVector &h, const EigenMatrix &J,
                       const EigenMatrix &qFixed, real G, real kT) {
    EigenMatrix dEmat = J * qFixed.transpose();
    real twoDivM = real(2.) / m_;
    real tempCoef = std::log(std::tanh(G / kT / m_)) / kT;
    real invKT = real(1.) / kT;

#ifndef _OPENMP
    {
        Random &random = random_[0];
#else
#pragma omp parallel
    {
        Random &random = random_[omp_get_thread_num()];
#pragma omp for
#endif
        for (int im = 0; im < m_; im += 2) {
            for (int iq = 0; iq < N; ++iq) {
                real q = qAnneal(im, iq);
                real dE = - twoDivM * q * (h[iq] + dEmat(iq, im));
                int mNeibour0 = (im + m_ - 1) % m_;
                int mNeibour1 = (im + 1) % m_;
                dE -= q * (qAnneal(mNeibour0, iq) + qAnneal(mNeibour1, iq)) * tempCoef;
                real thresh = dE < real(0.) ? real(1.) : std::exp(- dE * invKT);
                if (thresh > random.random<real>())
                    qAnneal(im, iq) = -q;
            }
        }
#ifdef _OPENMP
#pragma omp for
#endif
        for (int im = 1; im < m_; im += 2) {
            for (int iq = 0; iq < N; ++iq) {
                real q = qAnneal(im, iq);
                real dE = - twoDivM * q * (h[iq] + dEmat(iq, im));
                int mNeibour0 = (im + m_ - 1) % m_;
                int mNeibour1 = (im + 1) % m_;
                dE -= q * (qAnneal(mNeibour0, iq) + qAnneal(mNeibour1, iq)) * tempCoef;
                real thresh = dE < real(0.) ? real(1.) : std::exp(-dE * invKT);
                if (thresh > random.random<real>())
                    qAnneal(im, iq) = -q;
            }
        }
    }
}

template<class real>
void CPUBipartiteGraphAnnealer<real>::syncBits() {
    bitsPairX_.clear();
    bitsPairQ_.clear();
    Bits x0, x1;
    for (int idx = 0; idx < IdxType(m_); ++idx) {
        Bits q0 = extractRow<char>(matQ0_, idx);
        Bits q1 = extractRow<char>(matQ1_, idx);
        bitsPairQ_.pushBack(BitsPairArray::ValueType(q0, q1));
        Bits x0 = x_from_q(q0), x1 = x_from_q(q1);
        bitsPairX_.pushBack(BitsPairArray::ValueType(x0, x1));
    }
}

template class sqaod::CPUBipartiteGraphAnnealer<float>;
template class sqaod::CPUBipartiteGraphAnnealer<double>;
<commit_msg>Cleaning up, and checked with g++.<commit_after>#include "CPUBipartiteGraphAnnealer.h"
#include <cmath>
#include <float.h>
#include <algorithm>
#include <exception>
#include "CPUFormulas.h"

using namespace sqaod;

template<class real>
CPUBipartiteGraphAnnealer<real>::CPUBipartiteGraphAnnealer() {
    m_ = -1;
    annState_ = annNone;
    annealMethod_ = &CPUBipartiteGraphAnnealer::annealOneStepColoring;
#ifdef _OPENMP
    nProcs_ = omp_get_num_procs();
    log("# processors: %d", nProcs_);
#else
    nProcs_ = 1;
#endif
    random_ = new Random[nProcs_];
}

template<class real>
CPUBipartiteGraphAnnealer<real>::~CPUBipartiteGraphAnnealer() {
    delete [] random_;
}

template<class real>
void CPUBipartiteGraphAnnealer<real>::seed(unsigned long seed) {
    for (int idx = 0; idx < nProcs_; ++idx)
        random_[idx].seed(seed + 17 * idx);
    annState_ |= annRandSeedGiven;
}

template<class real>
void CPUBipartiteGraphAnnealer<real>::selectAlgorithm(Algorithm algo) {
    switch (algo) {
    case algoNaive:
        annealMethod_ = &CPUBipartiteGraphAnnealer::annealOneStepNaive;
        break;
    case algoColoring:
    case algoDefault:
        annealMethod_ = &CPUBipartiteGraphAnnealer::annealOneStepColoring;
        break;
    default:
        log("Uknown algo, %s, defaulting to %s.", algoToName(algo), algoToName(algoColoring));
        annealMethod_ = &CPUBipartiteGraphAnnealer::annealOneStepColoring;
    }
}

template<class real>
Algorithm CPUBipartiteGraphAnnealer<real>::algorithm() const {
    if (annealMethod_ == &CPUBipartiteGraphAnnealer::annealOneStepNaive)
        return algoNaive;
    if (annealMethod_ == &CPUBipartiteGraphAnnealer::annealOneStepColoring)
        return algoColoring;
    abort_("Must not reach here.");
    return algoDefault; /* to suppress warning. */
}

template<class real>
void CPUBipartiteGraphAnnealer<real>::getProblemSize(SizeType *N0, SizeType *N1) const {
    *N0 = N0_;
    *N1 = N1_;
}

template<class real>
void CPUBipartiteGraphAnnealer<real>::setProblem(const Vector &b0, const Vector &b1,
                                                 const Matrix &W, OptimizeMethod om) {
    N0_ = (int)b0.size;
    N1_ = (int)b1.size;
    h0_.resize(N0_);
    h1_.resize(N1_);
    J_.resize(N1_, N0_);
    Vector h0(mapFrom(h0_)), h1(mapFrom(h1_));
    Matrix J(mapFrom(J_));
    BGFuncs<real>::calculate_hJc(&h0, &h1, &J, &c_, b0, b1, W);

    om_ = om;
    if (om_ == optMaximize) {
        h0_ *= real(-1.);
        h1_ *= real(-1.);
        J_ *= real(-1.);
        c_ *= real(-1.);
    }
}

template<class real>
void CPUBipartiteGraphAnnealer<real>::setNumTrotters(SizeType m) {
    throwErrorIf(m <= 0, "# trotters must be a positive integer.");
    m_ = m;
    matQ0_.resize(m_, N0_);
    matQ1_.resize(m_, N1_);
    E_.resize(m_);
    annState_ |= annNTrottersGiven;
}

template<class real>
const BitsPairArray &CPUBipartiteGraphAnnealer<real>::get_x() const {
    return bitsPairX_;
}

template<class real>
void CPUBipartiteGraphAnnealer<real>::set_x(const Bits &x0, const Bits &x1) {
    EigenRowVector ex0 = mapToRowVector(x0).cast<real>();
    EigenRowVector ex1 = mapToRowVector(x1).cast<real>();
    matQ0_ = (ex0.array() * 2 - 1).matrix();
    matQ1_ = (ex1.array() * 2 - 1).matrix();
    annState_ |= annQSet;
}

template<class real>
const VectorType<real> &CPUBipartiteGraphAnnealer<real>::get_E() const {
    return E_;
}

template<class real>
void CPUBipartiteGraphAnnealer<real>::get_hJc(Vector *h0, Vector *h1,
                                              Matrix *J, real *c) const {
    mapToRowVector(*h0) = h0_;
    mapToRowVector(*h1) = h1_;
    mapTo(*J) = J_;
    *c = c_;
}

template<class real>
const BitsPairArray &CPUBipartiteGraphAnnealer<real>::get_q() const {
    return bitsPairQ_;
}

template<class real>
void CPUBipartiteGraphAnnealer<real>::randomize_q() {
#ifndef _OPENMP
    {
        Random &random = random_[0];
#else
#pragma omp parallel
    {
        sqaod::Random &random = random_[omp_get_thread_num()];
        real *q = matQ0_.data();
#pragma omp for
#endif
        for (int idx = 0; idx < IdxType(N0_ * m_); ++idx)
            q[idx] = random.randInt(2) ? real(1.) : real(-1.);
        q = matQ1_.data();
#ifdef _OPENMP
#pragma omp for
#endif
        for (int idx = 0; idx < IdxType(N1_ * m_); ++idx)
            q[idx] = random.randInt(2) ? real(1.) : real(-1.);
    }
    annState_ |= annQSet;
}

template<class real>
void CPUBipartiteGraphAnnealer<real>::calculate_E() {
    BGFuncs<real>::calculate_E(&E_, mapFrom(h0_), mapFrom(h1_), mapFrom(J_), c_,
                               mapFrom(matQ0_), mapFrom(matQ1_));
    if (om_ == optMaximize)
        mapToRowVector(E_) *= real(-1.);
}

template<class real>
void CPUBipartiteGraphAnnealer<real>::initAnneal() {
    if (!(annState_ & annRandSeedGiven))
        for (int idx = 0; idx < nProcs_; ++idx)
            random_[idx].seed();
    annState_ |= annRandSeedGiven;
    if (!(annState_ & annNTrottersGiven))
        setNumTrotters((N0_ + N1_) / 4);
    annState_ |= annNTrottersGiven;
    if (!(annState_ & annQSet))
        randomize_q();
    annState_ |= annQSet;
}

template<class real>
void CPUBipartiteGraphAnnealer<real>::finAnneal() {
    syncBits();
    calculate_E();
}

template<class real>
void CPUBipartiteGraphAnnealer<real>::annealOneStepNaive(real G, real kT) {
    real twoDivM = real(2.) / real(m_);
    real coef = std::log(std::tanh(G / kT / m_)) / kT;
    Random &random = random_[0];
    int N = N0_ + N1_;
    for (int loop = 0; loop < IdxType(N * m_); ++loop) {
        int x = random.randInt(N);
        int y = random.randInt(m_);
        if (x < N0_) {
            real qyx = matQ0_(y, x);
            real sum = J_.transpose().row(x).dot(matQ1_.row(y));
            real dE = - twoDivM * qyx * (h0_(x) + sum);
            int neibour0 = (y == 0) ? m_ - 1 : y - 1;
            int neibour1 = (y == m_ - 1) ? 0 : y + 1;
            dE -= qyx * (matQ0_(neibour0, x) + matQ0_(neibour1, x)) * coef;
            real threshold = (dE < real(0.)) ? real(1.) : std::exp(-dE / kT);
            if (threshold > random.random<real>())
                matQ0_(y, x) = - qyx;
        }
        else {
            x -= N0_;
            real qyx = matQ1_(y, x);
            real sum = J_.row(x).dot(matQ0_.row(y));
            real dE = - twoDivM * qyx * (h1_(x) + sum);
            int neibour0 = (y == 0) ? m_ - 1 : y - 1;
            int neibour1 = (y == m_ - 1) ? 0 : y + 1;
            dE -= qyx * (matQ1_(neibour0, x) + matQ1_(neibour1, x)) * coef;
            real threshold = (dE < real(0.)) ? real(1.) : std::exp(-dE / kT);
            if (threshold > random.random<real>())
                matQ1_(y, x) = - qyx;
        }
    }
}

template<class real>
void CPUBipartiteGraphAnnealer<real>::annealOneStepColoring(real G, real kT) {
    annealHalfStepColoring(N1_, matQ1_, h1_, J_, matQ0_, G, kT);
    annealHalfStepColoring(N0_, matQ0_, h0_, J_.transpose(), matQ1_, G, kT);
}

template<class real>
void CPUBipartiteGraphAnnealer<real>::
annealHalfStepColoring(int N, EigenMatrix &qAnneal,
                       const EigenRowVector &h, const EigenMatrix &J,
                       const EigenMatrix &qFixed, real G, real kT) {
    EigenMatrix dEmat = J * qFixed.transpose();
    real twoDivM = real(2.) / m_;
    real tempCoef = std::log(std::tanh(G / kT / m_)) / kT;
    real invKT = real(1.) / kT;

#ifndef _OPENMP
    {
        Random &random = random_[0];
#else
#pragma omp parallel
    {
        Random &random = random_[omp_get_thread_num()];
#pragma omp for
#endif
        for (int im = 0; im < m_; im += 2) {
            for (int iq = 0; iq < N; ++iq) {
                real q = qAnneal(im, iq);
                real dE = - twoDivM * q * (h[iq] + dEmat(iq, im));
                int mNeibour0 = (im + m_ - 1) % m_;
                int mNeibour1 = (im + 1) % m_;
                dE -= q * (qAnneal(mNeibour0, iq) + qAnneal(mNeibour1, iq)) * tempCoef;
                real thresh = dE < real(0.) ? real(1.) : std::exp(- dE * invKT);
                if (thresh > random.random<real>())
                    qAnneal(im, iq) = -q;
            }
        }
#ifdef _OPENMP
#pragma omp for
#endif
        for (int im = 1; im < m_; im += 2) {
            for (int iq = 0; iq < N; ++iq) {
                real q = qAnneal(im, iq);
                real dE = - twoDivM * q * (h[iq] + dEmat(iq, im));
                int mNeibour0 = (im + m_ - 1) % m_;
                int mNeibour1 = (im + 1) % m_;
                dE -= q * (qAnneal(mNeibour0, iq) + qAnneal(mNeibour1, iq)) * tempCoef;
                real thresh = dE < real(0.) ? real(1.) : std::exp(-dE * invKT);
                if (thresh > random.random<real>())
                    qAnneal(im, iq) = -q;
            }
        }
    }
}

template<class real>
void CPUBipartiteGraphAnnealer<real>::syncBits() {
    bitsPairX_.clear();
    bitsPairQ_.clear();
    Bits x0, x1;
    for (int idx = 0; idx < IdxType(m_); ++idx) {
        Bits q0 = extractRow<char>(matQ0_, idx);
        Bits q1 = extractRow<char>(matQ1_, idx);
        bitsPairQ_.pushBack(BitsPairArray::ValueType(q0, q1));
        Bits x0 = x_from_q(q0), x1 = x_from_q(q1);
        bitsPairX_.pushBack(BitsPairArray::ValueType(x0, x1));
    }
}

template class sqaod::CPUBipartiteGraphAnnealer<float>;
template class sqaod::CPUBipartiteGraphAnnealer<double>;
<|endoftext|>


<commit_before>// The MIT License (MIT)

// Copyright (c) 2016, Microsoft

// 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 "gtest/gtest.h"

#include "Allocator.h"
#include "AbstractRowEnumerator.h"
// #include "BitFunnel/Factories.h"
#include "BitFunnel/IFactSet.h"
#include "BitFunnel/ITermTable.h"
#include "BitFunnel/TermInfo.h"
// #include "MockIndexConfiguration.h"
// n#include "MockTermTable.h"
// #include "MockTermTableCollection.h"
#include "PlanRows.h"
#include "Random.h"

namespace BitFunnel
{
    namespace AbstractRowEnumeratorUnitTest
    {
        class RestrictedCapacityPlanRows : public PlanRows
        {
        public:

            RestrictedCapacityPlanRows(const IIndexConfiguration& index);

            virtual ~RestrictedCapacityPlanRows();

            const static unsigned c_maxRowsPerQuery = 5;

        protected:

            virtual unsigned GetRowCountLimit() const;
        };

        RestrictedCapacityPlanRows::RestrictedCapacityPlanRows(const IIndexConfiguration& index)
            : PlanRows(index)
        {
        }

        RestrictedCapacityPlanRows::~RestrictedCapacityPlanRows()
        {
        }

        unsigned RestrictedCapacityPlanRows::GetRowCountLimit() const
        {
            return c_maxRowsPerQuery;
        }

        static std::vector<DocIndex> CreateDefaultShardCapacities()
        {
            const std::vector<DocIndex> capacities = { 4096 };
            return capacities;
        }

        static std::vector<DocIndex> s_defaultShardCapacities(CreateDefaultShardCapacities());

        // Generate a random term list.
        void GenerateTermList(std::vector<Term>& termList, size_t count)
        {
            // Random term hash generator.
            RandomInt<Term::Hash> randomTermHashGenerator(7633134, 1, 4294967296);

            termList.clear();

            for (size_t i = 0; i < count; ++i)
            {
                const Term::Hash hash = randomTermHashGenerator();

                // hash, stream, idf.
                Term term(hash, 0, 10);

                termList.push_back(term);
            }
        }

        // Generate PlanRows for a list of terms.
        void GeneratePlanRows(IPlanRows& planRows, const std::vector<Term>& termList)
        {
            for (size_t i = 0; i < termList.size(); ++i)
            {
                AbstractRowEnumerator rowEnumerator(termList[i], planRows);
            }
        }

        // Verify planRows whose rows are truncated during the generation by AbstractRowEnumerator.
        void VerifyTruncatedPlanRows(size_t termCount)
        {
            // Test setup.
            MockIndexConfiguration index(s_defaultShardCapacities);
            PrivateHeapAllocator allocator;

            std::vector<Term> termList;

            GenerateTermList(termList, termCount);

            // First, generate a non-truncated plan row.
            PlanRows planRows(index);
            GeneratePlanRows(planRows, termList);

            // Then, generate a truncated plan row using the same
            // list of terms.
            RestrictedCapacityPlanRows truncatedPlanRows(index);
            GeneratePlanRows(truncatedPlanRows, termList);

            // Test the assumption of the test.
            TestEqual(static_cast<unsigned>(5),
                      RestrictedCapacityPlanRows::c_maxRowsPerQuery);

            // Compare the truncated plan row with the normal plan row.
            // Testing two conditions:
            // 1. The truncated plan row should have exactly RestrictedCapacityPlanRows::c_maxRowsPerQuery
            //    number of rows.
            // 2. The rows in the truncated plan rows should be in the normal plan rows as well.
            TestEqual(RestrictedCapacityPlanRows::c_maxRowsPerQuery,
                      truncatedPlanRows.GetRowCount());

            for (size_t rowIter = 0;
                 rowIter < RestrictedCapacityPlanRows::c_maxRowsPerQuery;
                 ++rowIter)
            {
                TestEqual(truncatedPlanRows.PhysicalRow(0, static_cast<unsigned>(rowIter)).GetShard(),
                          planRows.PhysicalRow(0, static_cast<unsigned>(rowIter)).GetShard());

                TestEqual(truncatedPlanRows.PhysicalRow(0, static_cast<unsigned>(rowIter)).GetTier(),
                          planRows.PhysicalRow(0, static_cast<unsigned>(rowIter)).GetTier());

                TestEqual(truncatedPlanRows.PhysicalRow(0, static_cast<unsigned>(rowIter)).GetRank(),
                          planRows.PhysicalRow(0, static_cast<unsigned>(rowIter)).GetRank());

                TestEqual(truncatedPlanRows.PhysicalRow(0, static_cast<unsigned>(rowIter)).GetIndex(),
                          planRows.PhysicalRow(0, static_cast<unsigned>(rowIter)).GetIndex());
            }
        }

        // Verify planRows generated by AbstractRowEnumerator.
        void VerifyPlanRows(size_t termCount)
        {
            MockIndexConfiguration index(s_defaultShardCapacities);
            PrivateHeapAllocator allocator;
            PlanRows planRows(index);

            std::vector<Term> termList;

            GenerateTermList(termList, termCount);

            GeneratePlanRows(planRows, termList);

            size_t rowIter = 0;

            for (size_t i = 0; i < termList.size(); ++i)
            {
                TermInfo rows(termList[i], planRows.GetTermTable(0));
                while (rows.MoveNext())
                {
                    RowId rowId = rows.Current();

                    // Verify that the rowIds are correctly set in the planRows.
                    TestEqual(rowId.GetShard(),
                              planRows.PhysicalRow(0, static_cast<unsigned>(rowIter)).GetShard());

                    TestEqual(rowId.GetTier(),
                              planRows.PhysicalRow(0, static_cast<unsigned>(rowIter)).GetTier());

                    TestEqual(rowId.GetRank(),
                              planRows.PhysicalRow(0, static_cast<unsigned>(rowIter)).GetRank());

                    TestEqual(rowId.GetIndex(),
                              planRows.PhysicalRow(0, static_cast<unsigned>(rowIter)).GetIndex());

                    rowIter++;
                }
            }
        }

        // Verify planRows which contains special rows during the generation by AbstractRowEnumerator.
        void VerifyPlanRowsWithSpecialRows(std::vector<Term> const & termList,
                                           RowId expectedSpecialRowId,
                                           unsigned expectedSpecialRowCount,
                                           ShardId shard,
                                           IIndexConfiguration& index)
        {
            // Generate plan rows.
            PlanRows planRows(index);
            GeneratePlanRows(planRows, termList);

            // Verify the generated plan rows has the expected number of special rows.
            unsigned actualSpecialRowCount = 0;
            for (size_t rowIter = 0;
                 rowIter < planRows.GetRowCount();
                 ++rowIter)
            {
                if (planRows.PhysicalRow(shard, static_cast<unsigned>(rowIter)) == expectedSpecialRowId)
                {
                    actualSpecialRowCount++;
                }
            }

            TestEqual(actualSpecialRowCount, expectedSpecialRowCount);
        }

        TestCase(SimpleCaseOneTerm)
        {
            VerifyPlanRows(1);
        }

        TestCase(MultipleTermsCase)
        {
            VerifyPlanRows(5);
        }

        TestCase(TruncatedPlanRows)
        {
            VerifyTruncatedPlanRows(RestrictedCapacityPlanRows::c_maxRowsPerQuery);
        }

        TestCase(MatchAllRowsInPlanRows)
        {
            // For shard 0, create a term table which assigns 2 rows for all terms in all ranks.
            const std::shared_ptr<MockTermTable> termTableShard0(new MockTermTable(0, 2, 2, 2));

            // For shard 1, create a term table which assigns 1 row for rank 0 and no rows for rank
            // 3 and 6.
            const std::shared_ptr<MockTermTable> termTableShard1(new MockTermTable(1, 1, 0, 0));

            // Set these two customized term table in the index configuration.
            const std::vector<DocIndex> c_shardCapacity(2, 1);
            std::unique_ptr<IFactSet> factSet(Factories::CreateFactSet());
            std::shared_ptr<const ITermTableCollection> termTableCollection(new MockTermTableCollection(c_shardCapacity, *factSet));
            static_cast<MockTermTableCollection*>(const_cast<ITermTableCollection*>(termTableCollection.get()))->SetTermTable(0, termTableShard0);
            static_cast<MockTermTableCollection*>(const_cast<ITermTableCollection*>(termTableCollection.get()))->SetTermTable(1, termTableShard1);
            MockIndexConfiguration index(termTableCollection);

            // Get the row Id of the match all rows from term table for shard 1 (Shard 1 will get the
            // match-all padding rows).
            TermInfo matchAllTerm(ITermTable::GetMatchAllTerm(), *(index.GetTermTables().GetTermTable(1)));
            LogAssertB(matchAllTerm.MoveNext());
            const RowId matchAllRowId = matchAllTerm.Current();

            const unsigned termCount = 1;
            std::vector<Term> termList;
            GenerateTermList(termList, termCount);

            // Since for shard 1, the term has no rows for rank 3 and 6, so there will be 2 match-all
            // padding rows for rank 3 and 6. In total, there will be 4 match-all rows.
            const unsigned expectedMatchAllRowCount = 4;

            // Verify there should be 4 match-all padding rows in shard 1.
            VerifyPlanRowsWithSpecialRows(termList, matchAllRowId, expectedMatchAllRowCount, 1, index);

            // Verify there should be no match-all padding rows in shard 0.
            VerifyPlanRowsWithSpecialRows(termList, matchAllRowId, 0, 0, index);
        }

        TestCase(MatchNoneRowsInPlanRows)
        {
            // Create a term which in the HDD tier to trigger the
            // generation of match none rows.
            Term termInHDDTier(1234, Stream::Full, 10, HDDTier);

            const std::vector<DocIndex> c_shardCapacity(1, 1);
            MockIndexConfiguration index(c_shardCapacity);

            // Get the row Id of the match none rows from term table.
            TermInfo matchNoneTerm(ITermTable::GetMatchNoneTerm(), *(index.GetTermTables().GetTermTable(0)));
            LogAssertB(matchNoneTerm.MoveNext());
            const RowId matchNoneRowId = matchNoneTerm.Current();

            // Get the expected number of rows for the special term in the term table.
            unsigned expectedMatchNoneRowCount = 0;
            ITermTableCollection const & termTableCollection = index.GetTermTables();
            std::shared_ptr<ITermTable const> termTable = termTableCollection.GetTermTable(0);
            TermInfo matchNoneTermInfo(termInHDDTier, *termTable);
            while (matchNoneTermInfo.MoveNext())
            {
                ++expectedMatchNoneRowCount;
            }

            const unsigned termCount = 4;
            std::vector<Term> termList;
            GenerateTermList(termList, termCount);

            termList.push_back(termInHDDTier);

            VerifyPlanRowsWithSpecialRows(termList, matchNoneRowId, expectedMatchNoneRowCount, 0, index);
        }
    }
}
<commit_msg>Remove AbstractRowEnumeratorTest.<commit_after><|endoftext|>


<commit_before>/*
  Copyright (c) 2014-2015 DataStax

  Licensed under the Apache License, Version 2.0 (the "License");
  you may not use this file except in compliance with the License.
  You may obtain a copy of the License at

  http://www.apache.org/licenses/LICENSE-2.0

  Unless required by applicable law or agreed to in writing, software
  distributed under the License is distributed on an "AS IS" BASIS,
  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  See the License for the specific language governing permissions and
  limitations under the License.
*/

#ifdef STAND_ALONE
#   define BOOST_TEST_MODULE cassandra
#endif

#include "test_utils.hpp"

#include "cassandra.h"
#include "external_types.hpp"
#include "timestamp_generator.hpp"
#include "get_time.hpp"

#include <boost/format.hpp>
#include <boost/test/unit_test.hpp>

struct TimestampsTest : public test_utils::SingleSessionTest {
  typedef test_utils::CassBatchPtr CassBatchPtr;
  typedef test_utils::CassFuturePtr CassFuturePtr;
  typedef test_utils::CassStatementPtr CassStatementPtr;

  TimestampsTest()
    : SingleSessionTest(1, 0) {
    test_utils::execute_query(session, str(boost::format(test_utils::CREATE_KEYSPACE_SIMPLE_FORMAT)
                                           % test_utils::SIMPLE_KEYSPACE % "1"));
    test_utils::execute_query(session, str(boost::format("USE %s") % test_utils::SIMPLE_KEYSPACE));
  }

  CassStatementPtr create_insert_statement(const std::string table_name) {
    std::string query(str(boost::format("INSERT INTO %s (key, value) VALUES (?, ?)") % table_name));
    return CassStatementPtr(cass_statement_new(query.c_str(), 2));
  }

  cass_int64_t get_timestamp(const std::string table_name, const std::string& key) {
    cass_int64_t timestamp;
    test_utils::CassResultPtr result;
    std::string query = str(boost::format("SELECT writetime(value) FROM %s WHERE key = '%s'") % table_name % key);
    test_utils::execute_query(session, query, &result);
    BOOST_REQUIRE(cass_result_row_count(result.get()) > 0);
    BOOST_REQUIRE(cass_result_column_count(result.get()) > 0);
    const CassRow* row = cass_result_first_row(result.get());
    const CassValue* writetime = cass_row_get_column(row, 0);
    cass_value_get_int64(writetime, &timestamp);
    return timestamp;
  }
};

class TestTimestampGenerator : public cass::TimestampGenerator {
public:
  TestTimestampGenerator(int64_t timestamp)
    : TimestampGenerator(SERVER_SIDE) // Doesn't matter for this test
    , timestamp_(timestamp) { }

  virtual int64_t next() { return timestamp_; }

private:
  int64_t timestamp_;
};

BOOST_FIXTURE_TEST_SUITE(timestamps, TimestampsTest)

/**
 * Set timestamp directly on statement and batch
 *
 * Verifies that the timestamp set on a statement/batch is sent to the server.
 *
 * @since 2.1.0
 * @jira_ticket CPP-266
 * @test_category timestamps
 * @cassandra_version 2.1.x
 */
BOOST_AUTO_TEST_CASE(statement_and_batch)
{
  std::string table_name("table_" + test_utils::generate_unique_str(uuid_gen));
  std::string create_table = "CREATE TABLE " + table_name + "(key text PRIMARY KEY, value text)";
  test_utils::execute_query(session, create_table.c_str());

  // Statement
  {
    std::string key(test_utils::generate_unique_str(uuid_gen));
    CassStatementPtr statement(create_insert_statement(table_name));
    cass_statement_bind_string(statement.get(), 0, key.c_str());
    cass_statement_bind_string(statement.get(), 1, key.c_str());

    // Set the timestamp to a known value
    cass_statement_set_timestamp(statement.get(), 1234);

    CassFuturePtr future(cass_session_execute(session, statement.get()));
    BOOST_REQUIRE(cass_future_error_code(future.get()) == CASS_OK);

    BOOST_CHECK_EQUAL(get_timestamp(table_name, key), 1234);
  }

  // Batch
  {
    CassBatchPtr batch(cass_batch_new(CASS_BATCH_TYPE_LOGGED));

    std::string key1(test_utils::generate_unique_str(uuid_gen));
    CassStatementPtr statement1(create_insert_statement(table_name));
    cass_statement_bind_string(statement1.get(), 0, key1.c_str());
    cass_statement_bind_string(statement1.get(), 1, key1.c_str());
    cass_batch_add_statement(batch.get(), statement1.get());

    std::string key2(test_utils::generate_unique_str(uuid_gen));
    CassStatementPtr statement2(create_insert_statement(table_name));
    cass_statement_bind_string(statement2.get(), 0, key2.c_str());
    cass_statement_bind_string(statement2.get(), 1, key2.c_str());
    cass_batch_add_statement(batch.get(), statement2.get());

    // Set the timestamp to a known value
    cass_batch_set_timestamp(batch.get(), 1234);

    CassFuturePtr future(cass_session_execute_batch(session, batch.get()));
    BOOST_REQUIRE(cass_future_error_code(future.get()) == CASS_OK);

    BOOST_CHECK_EQUAL(get_timestamp(table_name, key1), 1234);
    BOOST_CHECK_EQUAL(get_timestamp(table_name, key2), 1234);
  }
}

/**
 * Test timestamp generator
 *
 * Verifies that a timestamp generator is used when a statement's or batch's
 * timestamp is not set directly.
 *
 * @since 2.1.0
 * @jira_ticket CPP-266
 * @test_category timestamps
 * @cassandra_version 2.1.x
 */
BOOST_AUTO_TEST_CASE(generator)
{
  cass::SharedRefPtr<TestTimestampGenerator> gen(new TestTimestampGenerator(1234));

  close_session();
  cass_cluster_set_timestamp_gen(cluster, CassTimestampGen::to(gen.get()));
  create_session();

  test_utils::execute_query(session, str(boost::format("USE %s") % test_utils::SIMPLE_KEYSPACE));

  std::string table_name("table_" + test_utils::generate_unique_str(uuid_gen));
  std::string create_table = "CREATE TABLE " + table_name + "(key text PRIMARY KEY, value text)";
  test_utils::execute_query(session, create_table.c_str());

  // Statement
  {
    std::string table_name("table_" + test_utils::generate_unique_str(uuid_gen));
    std::string create_table = "CREATE TABLE " + table_name + "(key text PRIMARY KEY, value text)";
    test_utils::execute_query(session, create_table.c_str());

    std::string key(test_utils::generate_unique_str(uuid_gen));
    CassStatementPtr statement(create_insert_statement(table_name));
    cass_statement_bind_string(statement.get(), 0, key.c_str());
    cass_statement_bind_string(statement.get(), 1, key.c_str());

    CassFuturePtr future(cass_session_execute(session, statement.get()));
    BOOST_REQUIRE(cass_future_error_code(future.get()) == CASS_OK);

    BOOST_CHECK_EQUAL(get_timestamp(table_name, key), 1234);
  }

  // Batch
  {
    CassBatchPtr batch(cass_batch_new(CASS_BATCH_TYPE_LOGGED));

    std::string key1(test_utils::generate_unique_str(uuid_gen));
    CassStatementPtr statement1(create_insert_statement(table_name));
    cass_statement_bind_string(statement1.get(), 0, key1.c_str());
    cass_statement_bind_string(statement1.get(), 1, key1.c_str());
    cass_batch_add_statement(batch.get(), statement1.get());

    std::string key2(test_utils::generate_unique_str(uuid_gen));
    CassStatementPtr statement2(create_insert_statement(table_name));
    cass_statement_bind_string(statement2.get(), 0, key2.c_str());
    cass_statement_bind_string(statement2.get(), 1, key2.c_str());
    cass_batch_add_statement(batch.get(), statement2.get());

    CassFuturePtr future(cass_session_execute_batch(session, batch.get()));
    BOOST_REQUIRE(cass_future_error_code(future.get()) == CASS_OK);

    BOOST_CHECK_EQUAL(get_timestamp(table_name, key1), 1234);
    BOOST_CHECK_EQUAL(get_timestamp(table_name, key2), 1234);
  }
}

/**
 * Test the default timestamp generator.
 *
 * Verifies that the timestamp is set by the server when no timestamp
 * generator is set and the timestamp is not set directly on the statement.
 *
 * @since 2.1.0
 * @jira_ticket CPP-266
 * @test_category timestamps
 * @cassandra_version 2.1.x
 */
BOOST_AUTO_TEST_CASE(server_side)
{
  // Server-side is the default timestamp generator
  std::string table_name("table_" + test_utils::generate_unique_str(uuid_gen));
  std::string create_table = "CREATE TABLE " + table_name + "(key text PRIMARY KEY, value text)";
  test_utils::execute_query(session, create_table.c_str());

  std::string key(test_utils::generate_unique_str(uuid_gen));
  CassStatementPtr statement(create_insert_statement(table_name));
  cass_statement_bind_string(statement.get(), 0, key.c_str());
  cass_statement_bind_string(statement.get(), 1, key.c_str());

  uint64_t timestamp = cass::get_time_since_epoch_ms();
  CassFuturePtr future(cass_session_execute(session, statement.get()));
  BOOST_REQUIRE(cass_future_error_code(future.get()) == CASS_OK);

  BOOST_CHECK_GE(get_timestamp(table_name, key), static_cast<int64_t>(timestamp));
}

BOOST_AUTO_TEST_SUITE_END()

<commit_msg>Adding C* version check for 2.1+ protocol timestamp feature<commit_after>/*
  Copyright (c) 2014-2015 DataStax

  Licensed under the Apache License, Version 2.0 (the "License");
  you may not use this file except in compliance with the License.
  You may obtain a copy of the License at

  http://www.apache.org/licenses/LICENSE-2.0

  Unless required by applicable law or agreed to in writing, software
  distributed under the License is distributed on an "AS IS" BASIS,
  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  See the License for the specific language governing permissions and
  limitations under the License.
*/

#ifdef STAND_ALONE
#   define BOOST_TEST_MODULE cassandra
#endif

#include "test_utils.hpp"

#include "cassandra.h"
#include "external_types.hpp"
#include "timestamp_generator.hpp"
#include "get_time.hpp"

#include <boost/format.hpp>
#include <boost/test/unit_test.hpp>

struct TimestampsTest : public test_utils::SingleSessionTest {
  typedef test_utils::CassBatchPtr CassBatchPtr;
  typedef test_utils::CassFuturePtr CassFuturePtr;
  typedef test_utils::CassStatementPtr CassStatementPtr;

  TimestampsTest()
    : SingleSessionTest(1, 0) {
    test_utils::execute_query(session, str(boost::format(test_utils::CREATE_KEYSPACE_SIMPLE_FORMAT)
                                           % test_utils::SIMPLE_KEYSPACE % "1"));
    test_utils::execute_query(session, str(boost::format("USE %s") % test_utils::SIMPLE_KEYSPACE));
  }

  CassStatementPtr create_insert_statement(const std::string table_name) {
    std::string query(str(boost::format("INSERT INTO %s (key, value) VALUES (?, ?)") % table_name));
    return CassStatementPtr(cass_statement_new(query.c_str(), 2));
  }

  cass_int64_t get_timestamp(const std::string table_name, const std::string& key) {
    cass_int64_t timestamp;
    test_utils::CassResultPtr result;
    std::string query = str(boost::format("SELECT writetime(value) FROM %s WHERE key = '%s'") % table_name % key);
    test_utils::execute_query(session, query, &result);
    BOOST_REQUIRE(cass_result_row_count(result.get()) > 0);
    BOOST_REQUIRE(cass_result_column_count(result.get()) > 0);
    const CassRow* row = cass_result_first_row(result.get());
    const CassValue* writetime = cass_row_get_column(row, 0);
    cass_value_get_int64(writetime, &timestamp);
    return timestamp;
  }
};

class TestTimestampGenerator : public cass::TimestampGenerator {
public:
  TestTimestampGenerator(int64_t timestamp)
    : TimestampGenerator(SERVER_SIDE) // Doesn't matter for this test
    , timestamp_(timestamp) { }

  virtual int64_t next() { return timestamp_; }

private:
  int64_t timestamp_;
};

BOOST_AUTO_TEST_SUITE(timestamps)

/**
 * Set timestamp directly on statement and batch
 *
 * Verifies that the timestamp set on a statement/batch is sent to the server.
 *
 * @since 2.1.0
 * @jira_ticket CPP-266
 * @test_category queries:timestamp
 * @cassandra_version 2.1.x
 */
BOOST_AUTO_TEST_CASE(statement_and_batch)
{
  CassVersion version = test_utils::get_version();
  if ((version.major >= 2 && version.minor >= 1) || version.major > 2) {
    TimestampsTest tester;
    std::string table_name("table_" + test_utils::generate_unique_str(tester.uuid_gen));
    std::string create_table = "CREATE TABLE " + table_name + "(key text PRIMARY KEY, value text)";
    test_utils::execute_query(tester.session, create_table.c_str());

    // Statement
    {
      std::string key(test_utils::generate_unique_str(tester.uuid_gen));
      test_utils::CassStatementPtr statement(tester.create_insert_statement(table_name));
      cass_statement_bind_string(statement.get(), 0, key.c_str());
      cass_statement_bind_string(statement.get(), 1, key.c_str());

      // Set the timestamp to a known value
      cass_statement_set_timestamp(statement.get(), 1234);

      test_utils::CassFuturePtr future(cass_session_execute(tester.session, statement.get()));
      BOOST_REQUIRE(cass_future_error_code(future.get()) == CASS_OK);

      BOOST_CHECK_EQUAL(tester.get_timestamp(table_name, key), 1234);
    }

    // Batch
    {
      test_utils::CassBatchPtr batch(cass_batch_new(CASS_BATCH_TYPE_LOGGED));

      std::string key1(test_utils::generate_unique_str(tester.uuid_gen));
      test_utils::CassStatementPtr statement1(tester.create_insert_statement(table_name));
      cass_statement_bind_string(statement1.get(), 0, key1.c_str());
      cass_statement_bind_string(statement1.get(), 1, key1.c_str());
      cass_batch_add_statement(batch.get(), statement1.get());

      std::string key2(test_utils::generate_unique_str(tester.uuid_gen));
      test_utils::CassStatementPtr statement2(tester.create_insert_statement(table_name));
      cass_statement_bind_string(statement2.get(), 0, key2.c_str());
      cass_statement_bind_string(statement2.get(), 1, key2.c_str());
      cass_batch_add_statement(batch.get(), statement2.get());

      // Set the timestamp to a known value
      cass_batch_set_timestamp(batch.get(), 1234);

      test_utils::CassFuturePtr future(cass_session_execute_batch(tester.session, batch.get()));
      BOOST_REQUIRE(cass_future_error_code(future.get()) == CASS_OK);

      BOOST_CHECK_EQUAL(tester.get_timestamp(table_name, key1), 1234);
      BOOST_CHECK_EQUAL(tester.get_timestamp(table_name, key2), 1234);
    }
  } else {
    boost::unit_test::unit_test_log_t::instance().set_threshold_level(boost::unit_test::log_messages);
    BOOST_TEST_MESSAGE("Unsupported Test for Cassandra v" << version.to_string() << ": Skipping timestamps/statement_and_batch");
    BOOST_REQUIRE(true);
  }
}

/**
 * Test timestamp generator
 *
 * Verifies that a timestamp generator is used when a statement's or batch's
 * timestamp is not set directly.
 *
 * @since 2.1.0
 * @jira_ticket CPP-266
 * @test_category queries:timestamp
 * @cassandra_version 2.1.x
 */
BOOST_AUTO_TEST_CASE(generator)
{
  CassVersion version = test_utils::get_version();
  if ((version.major >= 2 && version.minor >= 1) || version.major > 2) {
    TimestampsTest tester;
    cass::SharedRefPtr<TestTimestampGenerator> gen(new TestTimestampGenerator(1234));

    tester.close_session();
    cass_cluster_set_timestamp_gen(tester.cluster, CassTimestampGen::to(gen.get()));
    tester.create_session();

    test_utils::execute_query(tester.session, str(boost::format("USE %s") % test_utils::SIMPLE_KEYSPACE));

    std::string table_name("table_" + test_utils::generate_unique_str(tester.uuid_gen));
    std::string create_table = "CREATE TABLE " + table_name + "(key text PRIMARY KEY, value text)";
    test_utils::execute_query(tester.session, create_table.c_str());

    // Statement
    {
      std::string table_name("table_" + test_utils::generate_unique_str(tester.uuid_gen));
      std::string create_table = "CREATE TABLE " + table_name + "(key text PRIMARY KEY, value text)";
      test_utils::execute_query(tester.session, create_table.c_str());

      std::string key(test_utils::generate_unique_str(tester.uuid_gen));
      test_utils::CassStatementPtr statement(tester.create_insert_statement(table_name));
      cass_statement_bind_string(statement.get(), 0, key.c_str());
      cass_statement_bind_string(statement.get(), 1, key.c_str());

      test_utils::CassFuturePtr future(cass_session_execute(tester.session, statement.get()));
      BOOST_REQUIRE(cass_future_error_code(future.get()) == CASS_OK);

      BOOST_CHECK_EQUAL(tester.get_timestamp(table_name, key), 1234);
    }

    // Batch
    {
      test_utils::CassBatchPtr batch(cass_batch_new(CASS_BATCH_TYPE_LOGGED));

      std::string key1(test_utils::generate_unique_str(tester.uuid_gen));
      test_utils::CassStatementPtr statement1(tester.create_insert_statement(table_name));
      cass_statement_bind_string(statement1.get(), 0, key1.c_str());
      cass_statement_bind_string(statement1.get(), 1, key1.c_str());
      cass_batch_add_statement(batch.get(), statement1.get());

      std::string key2(test_utils::generate_unique_str(tester.uuid_gen));
      test_utils::CassStatementPtr statement2(tester.create_insert_statement(table_name));
      cass_statement_bind_string(statement2.get(), 0, key2.c_str());
      cass_statement_bind_string(statement2.get(), 1, key2.c_str());
      cass_batch_add_statement(batch.get(), statement2.get());

      test_utils::CassFuturePtr future(cass_session_execute_batch(tester.session, batch.get()));
      BOOST_REQUIRE(cass_future_error_code(future.get()) == CASS_OK);

      BOOST_CHECK_EQUAL(tester.get_timestamp(table_name, key1), 1234);
      BOOST_CHECK_EQUAL(tester.get_timestamp(table_name, key2), 1234);
    }
  } else {
    boost::unit_test::unit_test_log_t::instance().set_threshold_level(boost::unit_test::log_messages);
    BOOST_TEST_MESSAGE("Unsupported Test for Cassandra v" << version.to_string() << ": Skipping timestamps/generator");
    BOOST_REQUIRE(true);
  }
}

/**
 * Test the default timestamp generator.
 *
 * Verifies that the timestamp is set by the server when no timestamp
 * generator is set and the timestamp is not set directly on the statement.
 *
 * @since 2.1.0
 * @jira_ticket CPP-266
 * @test_category queries:timestamp
 * @cassandra_version 2.1.x
 */
BOOST_AUTO_TEST_CASE(server_side)
{
  CassVersion version = test_utils::get_version();
  if ((version.major >= 2 && version.minor >= 1) || version.major > 2) {
    TimestampsTest tester;
    // Server-side is the default timestamp generator
    std::string table_name("table_" + test_utils::generate_unique_str(tester.uuid_gen));
    std::string create_table = "CREATE TABLE " + table_name + "(key text PRIMARY KEY, value text)";
    test_utils::execute_query(tester.session, create_table.c_str());

    std::string key(test_utils::generate_unique_str(tester.uuid_gen));
    test_utils::CassStatementPtr statement(tester.create_insert_statement(table_name));
    cass_statement_bind_string(statement.get(), 0, key.c_str());
    cass_statement_bind_string(statement.get(), 1, key.c_str());

    uint64_t timestamp = cass::get_time_since_epoch_ms();
    test_utils::CassFuturePtr future(cass_session_execute(tester.session, statement.get()));
    BOOST_REQUIRE(cass_future_error_code(future.get()) == CASS_OK);

    BOOST_CHECK_GE(tester.get_timestamp(table_name, key), static_cast<int64_t>(timestamp));
  } else {
    boost::unit_test::unit_test_log_t::instance().set_threshold_level(boost::unit_test::log_messages);
    BOOST_TEST_MESSAGE("Unsupported Test for Cassandra v" << version.to_string() << ": Skipping timestamps/generator");
    BOOST_REQUIRE(true);
  }
}

BOOST_AUTO_TEST_SUITE_END()
<|endoftext|>