<commit_before>////////////////////////////////////////////////////////////////////////////////
// task - a command line task list manager.
//
// Copyright 2006 - 2010, Paul Beckingham.
// All rights reserved.
//
// This program is free software; you can redistribute it and/or modify it under
// the terms of the GNU General Public License as published by the Free Software
// Foundation; either version 2 of the License, or (at your option) any later
// version.
//
// This program 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 General Public License for more
// details.
//
// You should have received a copy of the GNU General Public License along with
// this program; if not, write to the
//
//     Free Software Foundation, Inc.,
//     51 Franklin Street, Fifth Floor,
//     Boston, MA
//     02110-1301
//     USA
//
////////////////////////////////////////////////////////////////////////////////

#include <iostream>
#include "Context.h"
#include "Hooks.h"

extern Context context;

////////////////////////////////////////////////////////////////////////////////
Hook::Hook ()
: event ("")
, file ("")
, function ("")
{
}

////////////////////////////////////////////////////////////////////////////////
Hook::Hook (const std::string& e, const std::string& f, const std::string& fn)
: event (e)
, file (f)
, function (fn)
{
}

////////////////////////////////////////////////////////////////////////////////
Hook::Hook (const Hook& other)
{
  event = other.event;
  file = other.file;
  function = other.function;
}

////////////////////////////////////////////////////////////////////////////////
Hook& Hook::operator= (const Hook& other)
{
  if (this != &other)
  {
    event = other.event;
    file = other.file;
    function = other.function;
  }

  return *this;
}

////////////////////////////////////////////////////////////////////////////////
Hooks::Hooks ()
{
}

////////////////////////////////////////////////////////////////////////////////
Hooks::~Hooks ()
{
}

////////////////////////////////////////////////////////////////////////////////
void Hooks::initialize ()
{
#ifdef HAVE_LIBLUA
  api.initialize ();
#endif

  // TODO Enumerate all hooks, and tell API about the script files it must load
  //      in order to call them.  Note that API will perform a deferred read,
  //      which means that if it isn't called, a script will not be loaded.

  std::vector <std::string> vars;
  context.config.all (vars);

  std::vector <std::string>::iterator it;
  for (it = vars.begin (); it != vars.end (); ++it)
  {
    std::string type;
    std::string name;
    std::string value;

    // "<type>.<name>"
    Nibbler n (*it);
    if (n.getUntil ('.', type) &&
        type == "hook"         &&
        n.skip ('.')           &&
        n.getUntilEOS (name))
    {
      std::string value = context.config.get (*it);
      Nibbler n (value);

      // <path>:<function> [, ...]
      while (!n.depleted ())
      {
        std::string file;
        std::string function;
        if (n.getUntil (':', file) &&
            n.skip (':')           &&
            n.getUntil (',', function))
        {
          context.debug (std::string ("Event '") + name + "' hooked by " + file + ", function " + function);
          Hook h (name, Path::expand (file), function);
          all.push_back (h);

          (void) n.skip (',');
        }
        else
          throw std::string ("Malformed hook definition '") + *it + "'";
      }
    }
  }
}

////////////////////////////////////////////////////////////////////////////////
void Hooks::setTaskId (int id)
{
#ifdef HAVE_LIBLUA
  task_id = id;
#endif
}

////////////////////////////////////////////////////////////////////////////////
bool Hooks::trigger (const std::string& event)
{
#ifdef HAVE_LIBLUA
  std::vector <Hook>::iterator it;
  for (it = all.begin (); it != all.end (); ++it)
  {
    if (it->event == event)
    {
      bool rc = true;
      std::string type;
      if (eventType (event, type))
      {
        context.debug (std::string ("Event ") + event + " triggered");

        // Figure out where to get the calling-context info from.
             if (type == "program") rc = api.callProgramHook (it->file, it->function);
        else if (type == "list")    rc = api.callListHook    (it->file, it->function/*, tasks*/);
        else if (type == "task")    rc = api.callTaskHook    (it->file, it->function, task_id);
        else if (type == "field")   rc = api.callFieldHook   (it->file, it->function, "field", "value");
      }
      else
        throw std::string ("Unrecognized hook event '") + event + "'";

      // If any hook returns false, stop.
      if (!rc)
        return false;
    }
  }
#endif

  return true;
}

////////////////////////////////////////////////////////////////////////////////
bool Hooks::eventType (const std::string& event, std::string& type)
{
  if (event == "post-start"         ||
      event == "pre-exit"           ||
      event == "pre-debug"          || event == "post-debug"          ||
      event == "pre-header"         || event == "post-header"         ||
      event == "pre-footnote"       || event == "post-footnote"       ||
      event == "pre-output"         || event == "post-output"         ||
      event == "pre-dispatch"       || event == "post-dispatch"       ||
      event == "pre-gc"             || event == "post-gc"             ||
      event == "pre-undo"           || event == "post-undo"           ||
      event == "pre-file-lock"      || event == "post-file-lock"      ||
      event == "pre-add-command"    || event == "post-add-command"    ||
      event == "pre-delete-command" || event == "post-delete-command" ||
      event == "pre-info-command"   || event == "post-info-command")
  {
    type = "program";
    return true;
  }
  else if (event == "?")
  {
    type = "list";
    return true;
  }
  else if (event == "pre-tag"       || event == "post-tag"       ||
           event == "pre-detag"     || event == "post-detag"     ||
           event == "pre-delete"    || event == "post-delete"    ||
           event == "pre-completed" || event == "post-completed")
  {
    type = "task";
    return true;
  }
  else if (event == "?")
  {
    type = "field";
    return true;
  }

  return false;
}

////////////////////////////////////////////////////////////////////////////////
<commit_msg>Enhancement - Hooks<commit_after>////////////////////////////////////////////////////////////////////////////////
// task - a command line task list manager.
//
// Copyright 2006 - 2010, Paul Beckingham.
// All rights reserved.
//
// This program is free software; you can redistribute it and/or modify it under
// the terms of the GNU General Public License as published by the Free Software
// Foundation; either version 2 of the License, or (at your option) any later
// version.
//
// This program 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 General Public License for more
// details.
//
// You should have received a copy of the GNU General Public License along with
// this program; if not, write to the
//
//     Free Software Foundation, Inc.,
//     51 Franklin Street, Fifth Floor,
//     Boston, MA
//     02110-1301
//     USA
//
////////////////////////////////////////////////////////////////////////////////

#include <iostream>
#include "Context.h"
#include "Hooks.h"
#include "Timer.h"

extern Context context;

////////////////////////////////////////////////////////////////////////////////
Hook::Hook ()
: event ("")
, file ("")
, function ("")
{
}

////////////////////////////////////////////////////////////////////////////////
Hook::Hook (const std::string& e, const std::string& f, const std::string& fn)
: event (e)
, file (f)
, function (fn)
{
}

////////////////////////////////////////////////////////////////////////////////
Hook::Hook (const Hook& other)
{
  event = other.event;
  file = other.file;
  function = other.function;
}

////////////////////////////////////////////////////////////////////////////////
Hook& Hook::operator= (const Hook& other)
{
  if (this != &other)
  {
    event = other.event;
    file = other.file;
    function = other.function;
  }

  return *this;
}

////////////////////////////////////////////////////////////////////////////////
Hooks::Hooks ()
{
}

////////////////////////////////////////////////////////////////////////////////
Hooks::~Hooks ()
{
}

////////////////////////////////////////////////////////////////////////////////
// Enumerate all hooks, and tell API about the script files it must load in
// order to call them.  Note that API will perform a deferred read, which means
// that if it isn't called, a script will not be loaded.
void Hooks::initialize ()
{
#ifdef HAVE_LIBLUA
  api.initialize ();
#endif

  std::vector <std::string> vars;
  context.config.all (vars);

  std::vector <std::string>::iterator it;
  for (it = vars.begin (); it != vars.end (); ++it)
  {
    std::string type;
    std::string name;
    std::string value;

    // "<type>.<name>"
    Nibbler n (*it);
    if (n.getUntil ('.', type) &&
        type == "hook"         &&
        n.skip ('.')           &&
        n.getUntilEOS (name))
    {
      std::string value = context.config.get (*it);
      Nibbler n (value);

      // <path>:<function> [, ...]
      while (!n.depleted ())
      {
        std::string file;
        std::string function;
        if (n.getUntil (':', file) &&
            n.skip (':')           &&
            n.getUntil (',', function))
        {
          context.debug (std::string ("Event '") + name + "' hooked by " + file + ", function " + function);
          Hook h (name, Path::expand (file), function);
          all.push_back (h);

          (void) n.skip (',');
        }
        else
          throw std::string ("Malformed hook definition '") + *it + "'";
      }
    }
  }
}

////////////////////////////////////////////////////////////////////////////////
void Hooks::setTaskId (int id)
{
#ifdef HAVE_LIBLUA
  task_id = id;
#endif
}

////////////////////////////////////////////////////////////////////////////////
bool Hooks::trigger (const std::string& event)
{
#ifdef HAVE_LIBLUA
  std::vector <Hook>::iterator it;
  for (it = all.begin (); it != all.end (); ++it)
  {
    if (it->event == event)
    {
      Timer timer (std::string ("Hooks::trigger ") + event);

      bool rc = true;
      std::string type;
      if (eventType (event, type))
      {
        context.debug (std::string ("Event ") + event + " triggered");

        // Figure out where to get the calling-context info from.
             if (type == "program") rc = api.callProgramHook (it->file, it->function);
        else if (type == "list")    rc = api.callListHook    (it->file, it->function/*, tasks*/);
        else if (type == "task")    rc = api.callTaskHook    (it->file, it->function, task_id);
        else if (type == "field")   rc = api.callFieldHook   (it->file, it->function, "field", "value");
      }
      else
        throw std::string ("Unrecognized hook event '") + event + "'";

      // If any hook returns false, stop.
      if (!rc)
        return false;
    }
  }
#endif

  return true;
}

////////////////////////////////////////////////////////////////////////////////
bool Hooks::eventType (const std::string& event, std::string& type)
{
  if (event == "post-start"         ||
      event == "pre-exit"           ||
      event == "pre-debug"          || event == "post-debug"          ||
      event == "pre-header"         || event == "post-header"         ||
      event == "pre-footnote"       || event == "post-footnote"       ||
      event == "pre-output"         || event == "post-output"         ||
      event == "pre-dispatch"       || event == "post-dispatch"       ||
      event == "pre-gc"             || event == "post-gc"             ||
      event == "pre-undo"           || event == "post-undo"           ||
      event == "pre-file-lock"      || event == "post-file-lock"      ||
      event == "pre-add-command"    || event == "post-add-command"    ||
      event == "pre-delete-command" || event == "post-delete-command" ||
      event == "pre-info-command"   || event == "post-info-command")
  {
    type = "program";
    return true;
  }
  else if (event == "?")
  {
    type = "list";
    return true;
  }
  else if (event == "pre-tag"       || event == "post-tag"       ||
           event == "pre-detag"     || event == "post-detag"     ||
           event == "pre-delete"    || event == "post-delete"    ||
           event == "pre-completed" || event == "post-completed")
  {
    type = "task";
    return true;
  }
  else if (event == "?")
  {
    type = "field";
    return true;
  }

  return false;
}

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


<commit_before>/*
 * Layer.cpp
 *
 *  Created on: Jun 19, 2016
 *      Author: mad
 */

#include <vnl/ThreadEngine.h>
#include <vnl/Layer.h>
#include <vnl/Random.h>
#include <vnl/Pool.h>
#include <vnl/Process.h>
#include <vnl/ProcessClient.hxx>

#include <iostream>
#include <mutex>
#include <dirent.h>

namespace vnl {

uint64_t local_domain = 0;
const char* local_domain_name = 0;

volatile bool Layer::have_shutdown = false;
std::atomic<int> Layer::num_threads(0);

Map<String, String>* Layer::config = 0;

Random64* Random64::instance = 0;

Layer::Layer(const char* domain_name, const char* config_dir)
	:	closed(false)
{
	assert(Random64::instance == 0);
	assert(local_domain == 0);
	assert(global_pool == 0);
	assert(have_shutdown == false);
	assert(num_threads == 0);
	assert(Router::instance == 0);
	assert(Layer::config == 0);

	Random64::instance = new Random64();
	local_domain_name = domain_name;
	local_domain = vnl::hash64(domain_name);
	global_pool = new GlobalPool();
	config = new Map<String, String>();

	if(config_dir) {
		parse_config(config_dir);
	}

	Router::instance = new Router();

	vnl::spawn(new Process());
}

Layer::~Layer() {
	close();
}

void Layer::shutdown() {
	if(!have_shutdown) {
		ThreadEngine engine;
		ProcessClient proc = Address(local_domain, "vnl.Process");
		proc.connect(&engine);
		proc.shutdown();
		engine.flush();
		have_shutdown = true;
	}
}

void Layer::close() {
	if(closed) {
		return;
	}
	while(num_threads.load() > 0) {
		usleep(10*1000);
	}

	delete Router::instance;
	delete Random64::instance;
	delete global_pool;
	delete config;

	Page::cleanup();
	Block::cleanup();
	closed = true;
}

const String* Layer::get_config(const String& domain, const String& topic, const String& name) {
	static std::mutex mutex;
	String key;
	key << domain << ":" << topic << "->" << name;
	mutex.lock();
	String* value = config->find(key);
	mutex.unlock();
	return value;
}

void Layer::parse_config(const char* root_dir) {
	char buf[4096];
	// try to find parent
	{
		(String(root_dir) << "/parent").to_string(buf, sizeof(buf));
		FILE* parent = ::fopen(buf, "r");
		if(parent) {
			int num = ::fread(buf, 1, sizeof(buf)-1, parent);
			if(num > 1) {
				buf[num-1] = 0;
				(String(root_dir) << "/" << buf).to_string(buf, sizeof(buf));
				parse_config(buf);
			}
			::fclose(parent);
		}
	}
	// parse the whole tree
	DIR* dir = ::opendir(root_dir);
	if(!dir) {
		std::cout << "ERROR: could not open: " << root_dir << std::endl;
		return;
	}
	dirent* domain;
	while((domain = ::readdir(dir)) != 0) {
		if(domain->d_type == DT_DIR && domain->d_name[0] != '.') {
			(String(root_dir) << "/" << domain->d_name).to_string(buf, sizeof(buf));
			DIR* dir = ::opendir(buf);
			if(!dir) {
				std::cout << "WARNING: could not open: " << buf << std::endl;
				continue;
			}
			dirent* topic;
			while((topic = ::readdir(dir)) != 0) {
				if(topic->d_type == DT_DIR && topic->d_name[0] != '.') {
					(String(root_dir) << "/" << domain->d_name << "/" << topic->d_name).to_string(buf, sizeof(buf));
					DIR* dir = ::opendir(buf);
					if(!dir) {
						std::cout << "WARNING: could not open: " << buf << std::endl;
						continue;
					}
					dirent* name;
					while((name = ::readdir(dir)) != 0) {
						if(name->d_type == DT_REG && name->d_name[0] != '.') {
							(String(root_dir) << "/" << domain->d_name << "/" << topic->d_name << "/" << name->d_name).to_string(buf, sizeof(buf));
							FILE* file = ::fopen(buf, "r");
							if(!file) {
								std::cout << "WARNING: could not open: " << buf << std::endl;
								continue;
							}
							String key = String(domain->d_name) << ":" << topic->d_name << "->" << name->d_name;
							String& value = (*config)[key];
							value.clear();
							while(true) {
								int count = sizeof(buf)-1;
								int num = ::fread(buf, 1, count, file);
								if(num > 0) {
									if(::feof(file) && buf[num-1] == '\n') {
										num--;
									}
									buf[num] = 0;
									value << buf;
								} else {
									break;
								}
							}
							//std::cout << "[config] " << key << " = " << value << std::endl;
							::fclose(file);
						}
					}
					::closedir(dir);
				}
			}
			::closedir(dir);
		}
	}
	::closedir(dir);
}

void GlobalLogWriter::write(const String& str) {
	vnl::LogMsg* msg = vnl::LogMsg::create();
	msg->domain = node->my_domain;
	msg->topic = node->my_topic;
	msg->level = level;
	msg->msg = str;
	node->publish(msg, local_domain_name, "vnl.log");
}

} // vnl
<commit_msg>shutdown deadlock fix<commit_after>/*
 * Layer.cpp
 *
 *  Created on: Jun 19, 2016
 *      Author: mad
 */

#include <vnl/ThreadEngine.h>
#include <vnl/Layer.h>
#include <vnl/Random.h>
#include <vnl/Pool.h>
#include <vnl/Process.h>
#include <vnl/ProcessClient.hxx>

#include <iostream>
#include <mutex>
#include <dirent.h>

namespace vnl {

uint64_t local_domain = 0;
const char* local_domain_name = 0;

volatile bool Layer::have_shutdown = false;
std::atomic<int> Layer::num_threads(0);

Map<String, String>* Layer::config = 0;

Random64* Random64::instance = 0;

Layer::Layer(const char* domain_name, const char* config_dir)
	:	closed(false)
{
	assert(Random64::instance == 0);
	assert(local_domain == 0);
	assert(global_pool == 0);
	assert(have_shutdown == false);
	assert(num_threads == 0);
	assert(Router::instance == 0);
	assert(Layer::config == 0);

	Random64::instance = new Random64();
	local_domain_name = domain_name;
	local_domain = vnl::hash64(domain_name);
	global_pool = new GlobalPool();
	config = new Map<String, String>();

	if(config_dir) {
		parse_config(config_dir);
	}

	Router::instance = new Router();

	vnl::spawn(new Process());
}

Layer::~Layer() {
	close();
}

void Layer::shutdown() {
	if(!have_shutdown) {
		ThreadEngine engine;
		ProcessClient proc = Address(local_domain, "vnl.Process");
		proc.set_fail(true);
		proc.connect(&engine);
		try {
			proc.shutdown();
		} catch(...) {}
		engine.flush();
		have_shutdown = true;
	}
}

void Layer::close() {
	if(closed) {
		return;
	}
	while(num_threads.load() > 0) {
		usleep(10*1000);
	}

	delete Router::instance;
	delete Random64::instance;
	delete global_pool;
	delete config;

	Page::cleanup();
	Block::cleanup();
	closed = true;
}

const String* Layer::get_config(const String& domain, const String& topic, const String& name) {
	static std::mutex mutex;
	String key;
	key << domain << ":" << topic << "->" << name;
	mutex.lock();
	String* value = config->find(key);
	mutex.unlock();
	return value;
}

void Layer::parse_config(const char* root_dir) {
	char buf[4096];
	// try to find parent
	{
		(String(root_dir) << "/parent").to_string(buf, sizeof(buf));
		FILE* parent = ::fopen(buf, "r");
		if(parent) {
			int num = ::fread(buf, 1, sizeof(buf)-1, parent);
			if(num > 1) {
				buf[num-1] = 0;
				(String(root_dir) << "/" << buf).to_string(buf, sizeof(buf));
				parse_config(buf);
			}
			::fclose(parent);
		}
	}
	// parse the whole tree
	DIR* dir = ::opendir(root_dir);
	if(!dir) {
		std::cout << "ERROR: could not open: " << root_dir << std::endl;
		return;
	}
	dirent* domain;
	while((domain = ::readdir(dir)) != 0) {
		if(domain->d_type == DT_DIR && domain->d_name[0] != '.') {
			(String(root_dir) << "/" << domain->d_name).to_string(buf, sizeof(buf));
			DIR* dir = ::opendir(buf);
			if(!dir) {
				std::cout << "WARNING: could not open: " << buf << std::endl;
				continue;
			}
			dirent* topic;
			while((topic = ::readdir(dir)) != 0) {
				if(topic->d_type == DT_DIR && topic->d_name[0] != '.') {
					(String(root_dir) << "/" << domain->d_name << "/" << topic->d_name).to_string(buf, sizeof(buf));
					DIR* dir = ::opendir(buf);
					if(!dir) {
						std::cout << "WARNING: could not open: " << buf << std::endl;
						continue;
					}
					dirent* name;
					while((name = ::readdir(dir)) != 0) {
						if(name->d_type == DT_REG && name->d_name[0] != '.') {
							(String(root_dir) << "/" << domain->d_name << "/" << topic->d_name << "/" << name->d_name).to_string(buf, sizeof(buf));
							FILE* file = ::fopen(buf, "r");
							if(!file) {
								std::cout << "WARNING: could not open: " << buf << std::endl;
								continue;
							}
							String key = String(domain->d_name) << ":" << topic->d_name << "->" << name->d_name;
							String& value = (*config)[key];
							value.clear();
							while(true) {
								int count = sizeof(buf)-1;
								int num = ::fread(buf, 1, count, file);
								if(num > 0) {
									if(::feof(file) && buf[num-1] == '\n') {
										num--;
									}
									buf[num] = 0;
									value << buf;
								} else {
									break;
								}
							}
							//std::cout << "[config] " << key << " = " << value << std::endl;
							::fclose(file);
						}
					}
					::closedir(dir);
				}
			}
			::closedir(dir);
		}
	}
	::closedir(dir);
}

void GlobalLogWriter::write(const String& str) {
	vnl::LogMsg* msg = vnl::LogMsg::create();
	msg->domain = node->my_domain;
	msg->topic = node->my_topic;
	msg->level = level;
	msg->msg = str;
	node->publish(msg, local_domain_name, "vnl.log");
}

} // vnl
<|endoftext|>


<commit_before>// Scintilla source code edit control
/** @file LexVB.cxx
 ** Lexer for Visual Basic and VBScript.
 **/
// Copyright 1998-2003 by Neil Hodgson <neilh@scintilla.org>
// The License.txt file describes the conditions under which this software may be distributed.

#include <stdlib.h>
#include <string.h>
#include <ctype.h>
#include <stdio.h>
#include <stdarg.h>

#include "Platform.h"

#include "PropSet.h"
#include "Accessor.h"
#include "StyleContext.h"
#include "KeyWords.h"
#include "Scintilla.h"
#include "SciLexer.h"

static bool IsVBComment(Accessor &styler, int pos, int len) {
	return len>0 && styler[pos]=='\'';
}

static inline bool IsTypeCharacter(int ch) {
	return ch == '%' || ch == '&' || ch == '@' || ch == '!' || ch == '#' || ch == '$';
}

// Extended to accept accented characters
static inline bool IsAWordChar(int ch) {
	return ch >= 0x80 ||
	       (isalnum(ch) || ch == '.' || ch == '_');
}

static inline bool IsAWordStart(int ch) {
	return ch >= 0x80 ||
	       (isalnum(ch) || ch == '_');
}

static inline bool IsADateCharacter(const int ch) {
	return (ch < 0x80) &&
		(isalnum(ch) || ch == '|' || ch == '-' || ch == '/' || ch == ':' || ch == ' ' || ch == '\t');
}

static void ColouriseVBDoc(unsigned int startPos, int length, int initStyle,
                           WordList *keywordlists[], Accessor &styler, bool vbScriptSyntax) {

	WordList &keywords = *keywordlists[0];
	WordList &keywords2 = *keywordlists[1];
	WordList &keywords3 = *keywordlists[2];
	WordList &keywords4 = *keywordlists[3];

	styler.StartAt(startPos);

	int visibleChars = 0;

	StyleContext sc(startPos, length, initStyle, styler);

	for (; sc.More(); sc.Forward()) {

		if (sc.state == SCE_B_OPERATOR) {
			sc.SetState(SCE_B_DEFAULT);
		} else if (sc.state == SCE_B_IDENTIFIER) {
			if (!IsAWordChar(sc.ch)) {
				// In Basic (except VBScript), a variable name or a function name
				// can end with a special character indicating the type of the value
				// held or returned.
				bool skipType = false;
				if (!vbScriptSyntax && IsTypeCharacter(sc.ch)) {
					sc.Forward();	// Skip it
					skipType = true;
				}
				if (sc.ch == ']') {
					sc.Forward();
				}
				char s[100];
				sc.GetCurrentLowered(s, sizeof(s));
				if (skipType) {
					s[strlen(s) - 1] = '\0';
				}
				if (strcmp(s, "rem") == 0) {
					sc.ChangeState(SCE_B_COMMENT);
				} else {
					if (keywords.InList(s)) {
						sc.ChangeState(SCE_B_KEYWORD);
					} else if (keywords2.InList(s)) {
						sc.ChangeState(SCE_B_KEYWORD2);
					} else if (keywords3.InList(s)) {
						sc.ChangeState(SCE_B_KEYWORD3);
					} else if (keywords4.InList(s)) {
						sc.ChangeState(SCE_B_KEYWORD4);
					}	// Else, it is really an identifier...
					sc.SetState(SCE_B_DEFAULT);
				}
			}
		} else if (sc.state == SCE_B_NUMBER) {
			if (!IsAWordChar(sc.ch)) {
				sc.SetState(SCE_B_DEFAULT);
			}
		} else if (sc.state == SCE_B_STRING) {
			// VB doubles quotes to preserve them, so just end this string
			// state now as a following quote will start again
			if (sc.ch == '\"') {
				if (tolower(sc.chNext) == 'c') {
					sc.Forward();
				}
				sc.ForwardSetState(SCE_B_DEFAULT);
			} else if (sc.atLineEnd) {
				sc.ChangeState(SCE_B_STRINGEOL);
				sc.ForwardSetState(SCE_B_DEFAULT);
			}
		} else if (sc.state == SCE_B_COMMENT) {
			if (sc.atLineEnd) {
				sc.SetState(SCE_B_DEFAULT);
			}
		} else if (sc.state == SCE_B_PREPROCESSOR) {
			if (sc.atLineEnd) {
				sc.SetState(SCE_B_DEFAULT);
			}
		} else if (sc.state == SCE_B_DATE) {
			if (sc.ch == '#' || !IsADateCharacter(sc.chNext)) {
				sc.ForwardSetState(SCE_B_DEFAULT);
			}
		}

		if (sc.state == SCE_B_DEFAULT) {
			if (sc.ch == '\'') {
				sc.SetState(SCE_B_COMMENT);
			} else if (sc.ch == '\"') {
				sc.SetState(SCE_B_STRING);
			} else if (sc.ch == '#' && visibleChars == 0) {
				// Preprocessor commands are alone on their line
				sc.SetState(SCE_B_PREPROCESSOR);
			} else if (sc.ch == '#') {
				int n = 1;
				int chSeek = ' ';
				while ((n < 100) && (chSeek == ' ' || chSeek == '\t')) {
					chSeek = sc.GetRelative(n);
					n++;
				}
				if (IsADigit(chSeek)) {
					sc.SetState(SCE_B_DATE);
				} else {
					sc.SetState(SCE_B_OPERATOR);
				}
			} else if (sc.ch == '&' && tolower(sc.chNext) == 'h') {
				sc.SetState(SCE_B_NUMBER);
			} else if (sc.ch == '&' && tolower(sc.chNext) == 'o') {
				sc.SetState(SCE_B_NUMBER);
			} else if (IsADigit(sc.ch) || (sc.ch == '.' && IsADigit(sc.chNext))) {
				sc.SetState(SCE_B_NUMBER);
			} else if (IsAWordStart(sc.ch) || (sc.ch == '[')) {
				sc.SetState(SCE_B_IDENTIFIER);
			} else if (isoperator(static_cast<char>(sc.ch)) || (sc.ch == '\\')) {
				sc.SetState(SCE_B_OPERATOR);
			}
		}

		if (sc.atLineEnd) {
			visibleChars = 0;
		}
		if (!IsASpace(sc.ch)) {
			visibleChars++;
		}
	}
	sc.Complete();
}

static void FoldVBDoc(unsigned int startPos, int length, int,
						   WordList *[], Accessor &styler) {
	int endPos = startPos + length;

	// Backtrack to previous line in case need to fix its fold status
	int lineCurrent = styler.GetLine(startPos);
	if (startPos > 0) {
		if (lineCurrent > 0) {
			lineCurrent--;
			startPos = styler.LineStart(lineCurrent);
		}
	}
	int spaceFlags = 0;
	int indentCurrent = styler.IndentAmount(lineCurrent, &spaceFlags, IsVBComment);
	char chNext = styler[startPos];
	for (int i = startPos; i < endPos; i++) {
		char ch = chNext;
		chNext = styler.SafeGetCharAt(i + 1);

		if ((ch == '\r' && chNext != '\n') || (ch == '\n') || (i == endPos)) {
			int lev = indentCurrent;
			int indentNext = styler.IndentAmount(lineCurrent + 1, &spaceFlags, IsVBComment);
			if (!(indentCurrent & SC_FOLDLEVELWHITEFLAG)) {
				// Only non whitespace lines can be headers
				if ((indentCurrent & SC_FOLDLEVELNUMBERMASK) < (indentNext & SC_FOLDLEVELNUMBERMASK)) {
					lev |= SC_FOLDLEVELHEADERFLAG;
				} else if (indentNext & SC_FOLDLEVELWHITEFLAG) {
					// Line after is blank so check the next - maybe should continue further?
					int spaceFlags2 = 0;
					int indentNext2 = styler.IndentAmount(lineCurrent + 2, &spaceFlags2, IsVBComment);
					if ((indentCurrent & SC_FOLDLEVELNUMBERMASK) < (indentNext2 & SC_FOLDLEVELNUMBERMASK)) {
						lev |= SC_FOLDLEVELHEADERFLAG;
					}
				}
			}
			indentCurrent = indentNext;
			styler.SetLevel(lineCurrent, lev);
			lineCurrent++;
		}
	}
}

static void ColouriseVBNetDoc(unsigned int startPos, int length, int initStyle,
                           WordList *keywordlists[], Accessor &styler) {
	ColouriseVBDoc(startPos, length, initStyle, keywordlists, styler, false);
}

static void ColouriseVBScriptDoc(unsigned int startPos, int length, int initStyle,
                           WordList *keywordlists[], Accessor &styler) {
	ColouriseVBDoc(startPos, length, initStyle, keywordlists, styler, true);
}

static const char * const vbWordListDesc[] = {
	"Keywords",
	"user1",
	"user2",
	"user3",
	0
};

LexerModule lmVB(SCLEX_VB, ColouriseVBNetDoc, "vb", FoldVBDoc, vbWordListDesc);
LexerModule lmVBScript(SCLEX_VBSCRIPT, ColouriseVBScriptDoc, "vbscript", FoldVBDoc, vbWordListDesc);

<commit_msg>Patch from Philippe that improves handling of file numbers and date literals.<commit_after>// Scintilla source code edit control
/** @file LexVB.cxx
 ** Lexer for Visual Basic and VBScript.
 **/
// Copyright 1998-2005 by Neil Hodgson <neilh@scintilla.org>
// The License.txt file describes the conditions under which this software may be distributed.

#include <stdlib.h>
#include <string.h>
#include <ctype.h>
#include <stdio.h>
#include <stdarg.h>

#include "Platform.h"

#include "PropSet.h"
#include "Accessor.h"
#include "StyleContext.h"
#include "KeyWords.h"
#include "Scintilla.h"
#include "SciLexer.h"

// Internal state, highlighted as number
#define SCE_B_FILENUMBER SCE_B_DEFAULT+100

static bool IsVBComment(Accessor &styler, int pos, int len) {
	return len > 0 && styler[pos] == '\'';
}

static inline bool IsTypeCharacter(int ch) {
	return ch == '%' || ch == '&' || ch == '@' || ch == '!' || ch == '#' || ch == '$';
}

// Extended to accept accented characters
static inline bool IsAWordChar(int ch) {
	return ch >= 0x80 ||
	       (isalnum(ch) || ch == '.' || ch == '_');
}

static inline bool IsAWordStart(int ch) {
	return ch >= 0x80 ||
	       (isalpha(ch) || ch == '_');
}

static inline bool IsANumberChar(int ch) {
	// Not exactly following number definition (several dots are seen as OK, etc.)
	// but probably enough in most cases.
	return (ch < 0x80) &&
	        (isdigit(ch) || toupper(ch) == 'E' ||
             ch == '.' || ch == '-' || ch == '+');
}

static void ColouriseVBDoc(unsigned int startPos, int length, int initStyle,
                           WordList *keywordlists[], Accessor &styler, bool vbScriptSyntax) {

	WordList &keywords = *keywordlists[0];
	WordList &keywords2 = *keywordlists[1];
	WordList &keywords3 = *keywordlists[2];
	WordList &keywords4 = *keywordlists[3];

	styler.StartAt(startPos);

	int visibleChars = 0;
	int fileNbDigits = 0;

	// Do not leak onto next line
	if (initStyle == SCE_B_STRINGEOL || initStyle == SCE_B_COMMENT || initStyle == SCE_B_PREPROCESSOR) {
		initStyle = SCE_B_DEFAULT;
	}

	StyleContext sc(startPos, length, initStyle, styler);

	for (; sc.More(); sc.Forward()) {

		if (sc.state == SCE_B_OPERATOR) {
			sc.SetState(SCE_B_DEFAULT);
		} else if (sc.state == SCE_B_IDENTIFIER) {
			if (!IsAWordChar(sc.ch)) {
				// In Basic (except VBScript), a variable name or a function name
				// can end with a special character indicating the type of the value
				// held or returned.
				bool skipType = false;
				if (!vbScriptSyntax && IsTypeCharacter(sc.ch)) {
					sc.Forward();	// Skip it
					skipType = true;
				}
				if (sc.ch == ']') {
					sc.Forward();
				}
				char s[100];
				sc.GetCurrentLowered(s, sizeof(s));
				if (skipType) {
					s[strlen(s) - 1] = '\0';
				}
				if (strcmp(s, "rem") == 0) {
					sc.ChangeState(SCE_B_COMMENT);
				} else {
					if (keywords.InList(s)) {
						sc.ChangeState(SCE_B_KEYWORD);
					} else if (keywords2.InList(s)) {
						sc.ChangeState(SCE_B_KEYWORD2);
					} else if (keywords3.InList(s)) {
						sc.ChangeState(SCE_B_KEYWORD3);
					} else if (keywords4.InList(s)) {
						sc.ChangeState(SCE_B_KEYWORD4);
					}	// Else, it is really an identifier...
					sc.SetState(SCE_B_DEFAULT);
				}
			}
		} else if (sc.state == SCE_B_NUMBER) {
			// We stop the number definition on non-numerical non-dot non-eE non-sign char
			// Also accepts A-F for hex. numbers
			if (!IsANumberChar(sc.ch) && !(tolower(sc.ch) >= 'a' && tolower(sc.ch) <= 'f')) {
				sc.SetState(SCE_B_DEFAULT);
			}
		} else if (sc.state == SCE_B_STRING) {
			// VB doubles quotes to preserve them, so just end this string
			// state now as a following quote will start again
			if (sc.ch == '\"') {
				if (tolower(sc.chNext) == 'c') {
					sc.Forward();
				}
				sc.ForwardSetState(SCE_B_DEFAULT);
			} else if (sc.atLineEnd) {
				sc.ChangeState(SCE_B_STRINGEOL);
				sc.ForwardSetState(SCE_B_DEFAULT);
			}
		} else if (sc.state == SCE_B_COMMENT) {
			if (sc.atLineEnd) {
				sc.ForwardSetState(SCE_B_DEFAULT);
			}
		} else if (sc.state == SCE_B_PREPROCESSOR) {
			if (sc.atLineEnd) {
				sc.ForwardSetState(SCE_B_DEFAULT);
			}
		} else if (sc.state == SCE_B_FILENUMBER) {
			if (IsADigit(sc.ch)) {
				fileNbDigits++;
				if (fileNbDigits > 3) {
					sc.ChangeState(SCE_B_DATE);
				}
			} else if (sc.ch == '\r' || sc.ch == '\n' || sc.ch == ',') {
				// Regular uses: Close #1; Put #1, ...; Get #1, ... etc.
				// Too bad if date is format #27, Oct, 2003# or something like that...
				// Use regular number state
				sc.ChangeState(SCE_B_NUMBER);
				sc.SetState(SCE_B_DEFAULT);
			} else if (sc.ch == '#') {
				sc.ChangeState(SCE_B_DATE);
				sc.ForwardSetState(SCE_B_DEFAULT);
			} else {
				sc.ChangeState(SCE_B_DATE);
			}
			if (sc.state != SCE_B_FILENUMBER) {
				fileNbDigits = 0;
			}
		} else if (sc.state == SCE_B_DATE) {
			if (sc.atLineEnd) {
				sc.ChangeState(SCE_B_STRINGEOL);
				sc.ForwardSetState(SCE_B_DEFAULT);
			} else if (sc.ch == '#') {
				sc.ForwardSetState(SCE_B_DEFAULT);
			}
		}

		if (sc.state == SCE_B_DEFAULT) {
			if (sc.ch == '\'') {
				sc.SetState(SCE_B_COMMENT);
			} else if (sc.ch == '\"') {
				sc.SetState(SCE_B_STRING);
			} else if (sc.ch == '#' && visibleChars == 0) {
				// Preprocessor commands are alone on their line
				sc.SetState(SCE_B_PREPROCESSOR);
			} else if (sc.ch == '#') {
				// It can be a date literal, ending with #, or a file number, from 1 to 511
				// The date literal depends on the locale, so anything can go between #'s.
				// Can be #January 1, 1993# or #1 Jan 93# or #05/11/2003#, etc.
				// So we set the FILENUMBER state, and switch to DATE if it isn't a file number
				sc.SetState(SCE_B_FILENUMBER);
			} else if (sc.ch == '&' && tolower(sc.chNext) == 'h') {
				// Hexadecimal number
				sc.SetState(SCE_B_NUMBER);
				sc.Forward();
			} else if (sc.ch == '&' && tolower(sc.chNext) == 'o') {
				// Octal number
				sc.SetState(SCE_B_NUMBER);
				sc.Forward();
			} else if (IsADigit(sc.ch) || (sc.ch == '.' && IsADigit(sc.chNext))) {
				sc.SetState(SCE_B_NUMBER);
			} else if (IsAWordStart(sc.ch) || (sc.ch == '[')) {
				sc.SetState(SCE_B_IDENTIFIER);
			} else if (isoperator(static_cast<char>(sc.ch)) || (sc.ch == '\\')) {	// Integer division
				sc.SetState(SCE_B_OPERATOR);
			}
		}

		if (sc.atLineEnd) {
			visibleChars = 0;
		}
		if (!IsASpace(sc.ch)) {
			visibleChars++;
		}
	}
	sc.Complete();
}

static void FoldVBDoc(unsigned int startPos, int length, int,
						   WordList *[], Accessor &styler) {
	int endPos = startPos + length;

	// Backtrack to previous line in case need to fix its fold status
	int lineCurrent = styler.GetLine(startPos);
	if (startPos > 0) {
		if (lineCurrent > 0) {
			lineCurrent--;
			startPos = styler.LineStart(lineCurrent);
		}
	}
	int spaceFlags = 0;
	int indentCurrent = styler.IndentAmount(lineCurrent, &spaceFlags, IsVBComment);
	char chNext = styler[startPos];
	for (int i = startPos; i < endPos; i++) {
		char ch = chNext;
		chNext = styler.SafeGetCharAt(i + 1);

		if ((ch == '\r' && chNext != '\n') || (ch == '\n') || (i == endPos)) {
			int lev = indentCurrent;
			int indentNext = styler.IndentAmount(lineCurrent + 1, &spaceFlags, IsVBComment);
			if (!(indentCurrent & SC_FOLDLEVELWHITEFLAG)) {
				// Only non whitespace lines can be headers
				if ((indentCurrent & SC_FOLDLEVELNUMBERMASK) < (indentNext & SC_FOLDLEVELNUMBERMASK)) {
					lev |= SC_FOLDLEVELHEADERFLAG;
				} else if (indentNext & SC_FOLDLEVELWHITEFLAG) {
					// Line after is blank so check the next - maybe should continue further?
					int spaceFlags2 = 0;
					int indentNext2 = styler.IndentAmount(lineCurrent + 2, &spaceFlags2, IsVBComment);
					if ((indentCurrent & SC_FOLDLEVELNUMBERMASK) < (indentNext2 & SC_FOLDLEVELNUMBERMASK)) {
						lev |= SC_FOLDLEVELHEADERFLAG;
					}
				}
			}
			indentCurrent = indentNext;
			styler.SetLevel(lineCurrent, lev);
			lineCurrent++;
		}
	}
}

static void ColouriseVBNetDoc(unsigned int startPos, int length, int initStyle,
                           WordList *keywordlists[], Accessor &styler) {
	ColouriseVBDoc(startPos, length, initStyle, keywordlists, styler, false);
}

static void ColouriseVBScriptDoc(unsigned int startPos, int length, int initStyle,
                           WordList *keywordlists[], Accessor &styler) {
	ColouriseVBDoc(startPos, length, initStyle, keywordlists, styler, true);
}

static const char * const vbWordListDesc[] = {
	"Keywords",
	"user1",
	"user2",
	"user3",
	0
};

LexerModule lmVB(SCLEX_VB, ColouriseVBNetDoc, "vb", FoldVBDoc, vbWordListDesc);
LexerModule lmVBScript(SCLEX_VBSCRIPT, ColouriseVBScriptDoc, "vbscript", FoldVBDoc, vbWordListDesc);

<|endoftext|>


<commit_before>// Scintilla source code edit control
/** @file LexVB.cxx
 ** Lexer for Visual Basic and VBScript.
 **/
// Copyright 1998-2001 by Neil Hodgson <neilh@scintilla.org>
// The License.txt file describes the conditions under which this software may be distributed.

#include <stdlib.h>
#include <string.h>
#include <ctype.h>
#include <stdio.h>
#include <stdarg.h>

#include "Platform.h"

#include "PropSet.h"
#include "Accessor.h"
#include "StyleContext.h"
#include "KeyWords.h"
#include "Scintilla.h"
#include "SciLexer.h"

static bool IsVBComment(Accessor &styler, int pos, int len) {
	return len>0 && styler[pos]=='\'';
}

inline bool IsTypeCharacter(const int ch) {
	return ch == '%' || ch == '&' || ch == '@' || ch == '!' || ch == '#' || ch == '$';
}

static void ColouriseVBDoc(unsigned int startPos, int length, int initStyle,
                           WordList *keywordlists[], Accessor &styler, bool vbScriptSyntax) {

	WordList &keywords = *keywordlists[0];

	styler.StartAt(startPos);

	int visibleChars = 0;

	StyleContext sc(startPos, length, initStyle, styler);

	for (; sc.More(); sc.Forward()) {

		if (sc.state == SCE_B_OPERATOR) {
			sc.SetState(SCE_B_DEFAULT);
		} else if (sc.state == SCE_B_KEYWORD) {
			if (!iswordchar(sc.ch)) {
				if (vbScriptSyntax || !IsTypeCharacter(sc.ch)) {
					if (sc.ch == ']')
						sc.Forward();
					char s[100];
					sc.GetCurrentLowered(s, sizeof(s));
					if (keywords.InList(s)) {
						if (strcmp(s, "rem") == 0) {
							sc.ChangeState(SCE_C_COMMENTLINE);
						}
					} else {
						sc.ChangeState(SCE_B_IDENTIFIER);
					}
					sc.SetState(SCE_B_DEFAULT);
				}
			}
		} else if (sc.state == SCE_B_NUMBER) {
			if (!iswordchar(sc.ch)) {
				sc.SetState(SCE_B_DEFAULT);
			}
		} else if (sc.state == SCE_B_STRING) {
			// VB doubles quotes to preserve them, so just end this string
			// state now as a following quote will start again
			if (sc.ch == '\"') {
				if (tolower(sc.chNext) == 'c') {
					sc.Forward();
				}
				sc.ForwardSetState(SCE_B_DEFAULT);
			}
		} else if (sc.state == SCE_B_COMMENT) {
			if (sc.atLineEnd) {
				sc.SetState(SCE_B_DEFAULT);
			}
		} else if (sc.state == SCE_B_PREPROCESSOR) {
			if (sc.atLineEnd) {
				sc.SetState(SCE_B_DEFAULT);
			}
		} else if (sc.state == SCE_B_DATE) {
			if (sc.ch == '#') {
				sc.ForwardSetState(SCE_B_DEFAULT);
			}
		}

		if (sc.state == SCE_B_DEFAULT) {
			if (sc.ch == '\'') {
				sc.SetState(SCE_B_COMMENT);
			} else if (sc.ch == '\"') {
				sc.SetState(SCE_B_STRING);
			} else if (sc.ch == '#' && visibleChars == 0) {
				// Preprocessor commands are alone on their line
				sc.SetState(SCE_B_PREPROCESSOR);
			} else if (sc.ch == '#') {
				int n = 1;
				int chSeek = ' ';
				while (chSeek == ' ' || chSeek == '\t') {
					chSeek = sc.GetRelative(n);
					n++;
				}
				if (IsADigit(chSeek)) {
					sc.SetState(SCE_B_DATE);
				} else {
					sc.SetState(SCE_B_OPERATOR);
				}
			} else if (sc.ch == '&' && tolower(sc.chNext) == 'h') {
				sc.SetState(SCE_B_NUMBER);
			} else if (sc.ch == '&' && tolower(sc.chNext) == 'o') {
				sc.SetState(SCE_B_NUMBER);
			} else if (IsADigit(sc.ch) || (sc.ch == '.' && IsADigit(sc.chNext))) {
				sc.SetState(SCE_B_NUMBER);
			} else if (iswordstart(sc.ch) || (sc.ch == '[')) {
				sc.SetState(SCE_B_KEYWORD);
			} else if (isoperator(sc.ch) || (sc.ch == '\\')) {
				sc.SetState(SCE_B_OPERATOR);
			}
		}

		if (sc.atLineEnd) {
			visibleChars = 0;
		}
		if (!IsASpace(sc.ch)) {
			visibleChars++;
		}
	}
	sc.Complete();
}

static void FoldVBDoc(unsigned int startPos, int length, int initStyle,
						   WordList *[], Accessor &styler) {
	int endPos = startPos + length;

	// Backtrack to previous line in case need to fix its fold status
	int lineCurrent = styler.GetLine(startPos);
	if (startPos > 0) {
		if (lineCurrent > 0) {
			lineCurrent--;
			startPos = styler.LineStart(lineCurrent);
			if (startPos == 0)
				initStyle = SCE_B_DEFAULT;
			else
				initStyle = styler.StyleAt(startPos-1);
		}
	}
	int spaceFlags = 0;
	int indentCurrent = styler.IndentAmount(lineCurrent, &spaceFlags, IsVBComment);
	char chNext = styler[startPos];
	for (int i = startPos; i < endPos; i++) {
		char ch = chNext;
		chNext = styler.SafeGetCharAt(i + 1);

		if ((ch == '\r' && chNext != '\n') || (ch == '\n') || (i == endPos)) {
			int lev = indentCurrent;
			int indentNext = styler.IndentAmount(lineCurrent + 1, &spaceFlags, IsVBComment);
			if (!(indentCurrent & SC_FOLDLEVELWHITEFLAG)) {
				// Only non whitespace lines can be headers
				if ((indentCurrent & SC_FOLDLEVELNUMBERMASK) < (indentNext & SC_FOLDLEVELNUMBERMASK)) {
					lev |= SC_FOLDLEVELHEADERFLAG;
				} else if (indentNext & SC_FOLDLEVELWHITEFLAG) {
					// Line after is blank so check the next - maybe should continue further?
					int spaceFlags2 = 0;
					int indentNext2 = styler.IndentAmount(lineCurrent + 2, &spaceFlags2, IsVBComment);
					if ((indentCurrent & SC_FOLDLEVELNUMBERMASK) < (indentNext2 & SC_FOLDLEVELNUMBERMASK)) {
						lev |= SC_FOLDLEVELHEADERFLAG;
					}
				}
			}
			indentCurrent = indentNext;
			styler.SetLevel(lineCurrent, lev);
			lineCurrent++;
		}
	}
}

static void ColouriseVBNetDoc(unsigned int startPos, int length, int initStyle,
                           WordList *keywordlists[], Accessor &styler) {
	ColouriseVBDoc(startPos, length, initStyle, keywordlists, styler, false);
}

static void ColouriseVBScriptDoc(unsigned int startPos, int length, int initStyle,
                           WordList *keywordlists[], Accessor &styler) {
	ColouriseVBDoc(startPos, length, initStyle, keywordlists, styler, true);
}

LexerModule lmVB(SCLEX_VB, ColouriseVBNetDoc, "vb", FoldVBDoc);
LexerModule lmVBScript(SCLEX_VBSCRIPT, ColouriseVBScriptDoc, "vbscript", FoldVBDoc);
<commit_msg>Fixed comments that use rem rather than '.<commit_after>// Scintilla source code edit control
/** @file LexVB.cxx
 ** Lexer for Visual Basic and VBScript.
 **/
// Copyright 1998-2001 by Neil Hodgson <neilh@scintilla.org>
// The License.txt file describes the conditions under which this software may be distributed.

#include <stdlib.h>
#include <string.h>
#include <ctype.h>
#include <stdio.h>
#include <stdarg.h>

#include "Platform.h"

#include "PropSet.h"
#include "Accessor.h"
#include "StyleContext.h"
#include "KeyWords.h"
#include "Scintilla.h"
#include "SciLexer.h"

static bool IsVBComment(Accessor &styler, int pos, int len) {
	return len>0 && styler[pos]=='\'';
}

inline bool IsTypeCharacter(const int ch) {
	return ch == '%' || ch == '&' || ch == '@' || ch == '!' || ch == '#' || ch == '$';
}

static void ColouriseVBDoc(unsigned int startPos, int length, int initStyle,
                           WordList *keywordlists[], Accessor &styler, bool vbScriptSyntax) {

	WordList &keywords = *keywordlists[0];

	styler.StartAt(startPos);

	int visibleChars = 0;

	StyleContext sc(startPos, length, initStyle, styler);

	for (; sc.More(); sc.Forward()) {

		if (sc.state == SCE_B_OPERATOR) {
			sc.SetState(SCE_B_DEFAULT);
		} else if (sc.state == SCE_B_KEYWORD) {
			if (!iswordchar(sc.ch)) {
				if (vbScriptSyntax || !IsTypeCharacter(sc.ch)) {
					if (sc.ch == ']')
						sc.Forward();
					char s[100];
					sc.GetCurrentLowered(s, sizeof(s));
					if (keywords.InList(s)) {
						if (strcmp(s, "rem") == 0) {
							sc.ChangeState(SCE_B_COMMENT);
							if (sc.atLineEnd) {
								sc.SetState(SCE_B_DEFAULT);
							}
						} else {
							sc.SetState(SCE_B_DEFAULT);
						}
					} else {
						sc.ChangeState(SCE_B_IDENTIFIER);
						sc.SetState(SCE_B_DEFAULT);
					}
				}
			}
		} else if (sc.state == SCE_B_NUMBER) {
			if (!iswordchar(sc.ch)) {
				sc.SetState(SCE_B_DEFAULT);
			}
		} else if (sc.state == SCE_B_STRING) {
			// VB doubles quotes to preserve them, so just end this string
			// state now as a following quote will start again
			if (sc.ch == '\"') {
				if (tolower(sc.chNext) == 'c') {
					sc.Forward();
				}
				sc.ForwardSetState(SCE_B_DEFAULT);
			}
		} else if (sc.state == SCE_B_COMMENT) {
			if (sc.atLineEnd) {
				sc.SetState(SCE_B_DEFAULT);
			}
		} else if (sc.state == SCE_B_PREPROCESSOR) {
			if (sc.atLineEnd) {
				sc.SetState(SCE_B_DEFAULT);
			}
		} else if (sc.state == SCE_B_DATE) {
			if (sc.ch == '#') {
				sc.ForwardSetState(SCE_B_DEFAULT);
			}
		}

		if (sc.state == SCE_B_DEFAULT) {
			if (sc.ch == '\'') {
				sc.SetState(SCE_B_COMMENT);
			} else if (sc.ch == '\"') {
				sc.SetState(SCE_B_STRING);
			} else if (sc.ch == '#' && visibleChars == 0) {
				// Preprocessor commands are alone on their line
				sc.SetState(SCE_B_PREPROCESSOR);
			} else if (sc.ch == '#') {
				int n = 1;
				int chSeek = ' ';
				while (chSeek == ' ' || chSeek == '\t') {
					chSeek = sc.GetRelative(n);
					n++;
				}
				if (IsADigit(chSeek)) {
					sc.SetState(SCE_B_DATE);
				} else {
					sc.SetState(SCE_B_OPERATOR);
				}
			} else if (sc.ch == '&' && tolower(sc.chNext) == 'h') {
				sc.SetState(SCE_B_NUMBER);
			} else if (sc.ch == '&' && tolower(sc.chNext) == 'o') {
				sc.SetState(SCE_B_NUMBER);
			} else if (IsADigit(sc.ch) || (sc.ch == '.' && IsADigit(sc.chNext))) {
				sc.SetState(SCE_B_NUMBER);
			} else if (iswordstart(sc.ch) || (sc.ch == '[')) {
				sc.SetState(SCE_B_KEYWORD);
			} else if (isoperator(sc.ch) || (sc.ch == '\\')) {
				sc.SetState(SCE_B_OPERATOR);
			}
		}

		if (sc.atLineEnd) {
			visibleChars = 0;
		}
		if (!IsASpace(sc.ch)) {
			visibleChars++;
		}
	}
	sc.Complete();
}

static void FoldVBDoc(unsigned int startPos, int length, int initStyle,
						   WordList *[], Accessor &styler) {
	int endPos = startPos + length;

	// Backtrack to previous line in case need to fix its fold status
	int lineCurrent = styler.GetLine(startPos);
	if (startPos > 0) {
		if (lineCurrent > 0) {
			lineCurrent--;
			startPos = styler.LineStart(lineCurrent);
			if (startPos == 0)
				initStyle = SCE_B_DEFAULT;
			else
				initStyle = styler.StyleAt(startPos-1);
		}
	}
	int spaceFlags = 0;
	int indentCurrent = styler.IndentAmount(lineCurrent, &spaceFlags, IsVBComment);
	char chNext = styler[startPos];
	for (int i = startPos; i < endPos; i++) {
		char ch = chNext;
		chNext = styler.SafeGetCharAt(i + 1);

		if ((ch == '\r' && chNext != '\n') || (ch == '\n') || (i == endPos)) {
			int lev = indentCurrent;
			int indentNext = styler.IndentAmount(lineCurrent + 1, &spaceFlags, IsVBComment);
			if (!(indentCurrent & SC_FOLDLEVELWHITEFLAG)) {
				// Only non whitespace lines can be headers
				if ((indentCurrent & SC_FOLDLEVELNUMBERMASK) < (indentNext & SC_FOLDLEVELNUMBERMASK)) {
					lev |= SC_FOLDLEVELHEADERFLAG;
				} else if (indentNext & SC_FOLDLEVELWHITEFLAG) {
					// Line after is blank so check the next - maybe should continue further?
					int spaceFlags2 = 0;
					int indentNext2 = styler.IndentAmount(lineCurrent + 2, &spaceFlags2, IsVBComment);
					if ((indentCurrent & SC_FOLDLEVELNUMBERMASK) < (indentNext2 & SC_FOLDLEVELNUMBERMASK)) {
						lev |= SC_FOLDLEVELHEADERFLAG;
					}
				}
			}
			indentCurrent = indentNext;
			styler.SetLevel(lineCurrent, lev);
			lineCurrent++;
		}
	}
}

static void ColouriseVBNetDoc(unsigned int startPos, int length, int initStyle,
                           WordList *keywordlists[], Accessor &styler) {
	ColouriseVBDoc(startPos, length, initStyle, keywordlists, styler, false);
}

static void ColouriseVBScriptDoc(unsigned int startPos, int length, int initStyle,
                           WordList *keywordlists[], Accessor &styler) {
	ColouriseVBDoc(startPos, length, initStyle, keywordlists, styler, true);
}

LexerModule lmVB(SCLEX_VB, ColouriseVBNetDoc, "vb", FoldVBDoc);
LexerModule lmVBScript(SCLEX_VBSCRIPT, ColouriseVBScriptDoc, "vbscript", FoldVBDoc);
<|endoftext|>


<commit_before>/*****************************************************************************
 * Media Library
 *****************************************************************************
 * Copyright (C) 2015 Hugo Beauzée-Luyssen, Videolabs
 *
 * Authors: Hugo Beauzée-Luyssen<hugo@beauzee.fr>
 *
 * This program 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.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
 * GNU Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public License
 * along with this program; if not, write to the Free Software Foundation,
 * Inc., 51 Franklin Street, Fifth Floor, Boston MA 02110-1301, USA.
 *****************************************************************************/

#include <algorithm>
#include <cassert>
#include <cstdlib>
#include <cstring>
#include <ctime>

#include "Album.h"
#include "AlbumTrack.h"
#include "Artist.h"
#include "AudioTrack.h"
#include "Media.h"
#include "File.h"
#include "Folder.h"
#include "Label.h"
#include "logging/Logger.h"
#include "Movie.h"
#include "ShowEpisode.h"
#include "database/SqliteTools.h"
#include "VideoTrack.h"
#include "filesystem/IFile.h"
#include "filesystem/IDirectory.h"
#include "filesystem/IDevice.h"
#include "utils/Filename.h"

const std::string policy::MediaTable::Name = "Media";
const std::string policy::MediaTable::PrimaryKeyColumn = "id_media";
unsigned int Media::* const policy::MediaTable::PrimaryKey = &Media::m_id;

Media::Media( DBConnection dbConnection, sqlite::Row& row )
    : m_dbConnection( dbConnection )
    , m_changed( false )
{
    row >> m_id
        >> m_type
        >> m_subType
        >> m_duration
        >> m_playCount
        >> m_progress
        >> m_rating
        >> m_insertionDate
        >> m_thumbnail
        >> m_title;
}

Media::Media( const std::string& title, Type type )
    : m_id( 0 )
    , m_type( type )
    , m_subType( SubType::Unknown )
    , m_duration( -1 )
    , m_playCount( 0 )
    , m_progress( .0f )
    , m_rating( -1 )
    , m_insertionDate( time( nullptr ) )
    , m_title( title )
    , m_changed( false )
{
}

std::shared_ptr<Media> Media::create( DBConnection dbConnection, Type type, const fs::IFile* file )
{
    auto self = std::make_shared<Media>( file->name(), type );
    static const std::string req = "INSERT INTO " + policy::MediaTable::Name +
            "(type, insertion_date, title) VALUES(?, ?, ?)";

    if ( insert( dbConnection, self, req, type, self->m_insertionDate, self->m_title ) == false )
        return nullptr;
    self->m_dbConnection = dbConnection;
    return self;
}

AlbumTrackPtr Media::albumTrack() const
{
    if ( m_subType != SubType::AlbumTrack )
        return nullptr;
    auto lock = m_albumTrack.lock();

    if ( m_albumTrack.isCached() == false )
        m_albumTrack = AlbumTrack::fromMedia( m_dbConnection, m_id );
    return m_albumTrack.get();
}

void Media::setAlbumTrack( AlbumTrackPtr albumTrack )
{
    auto lock = m_albumTrack.lock();
    m_albumTrack = albumTrack;
    m_subType = SubType::AlbumTrack;
    m_changed = true;
}

int64_t Media::duration() const
{
    return m_duration;
}

void Media::setDuration( int64_t duration )
{
    if ( m_duration == duration )
        return;
    m_duration = duration;
    m_changed = true;
}

ShowEpisodePtr Media::showEpisode() const
{
    if ( m_subType != SubType::ShowEpisode )
        return nullptr;

    auto lock = m_showEpisode.lock();
    if ( m_showEpisode.isCached() == false )
        m_showEpisode = ShowEpisode::fromMedia( m_dbConnection, m_id );
    return m_showEpisode.get();
}

void Media::setShowEpisode( ShowEpisodePtr episode )
{
    auto lock = m_showEpisode.lock();
    m_showEpisode = episode;
    m_subType = SubType::ShowEpisode;
    m_changed = true;
}

std::vector<LabelPtr> Media::labels()
{
    static const std::string req = "SELECT l.* FROM " + policy::LabelTable::Name + " l "
            "INNER JOIN LabelFileRelation lfr ON lfr.label_id = l.id_label "
            "WHERE lfr.media_id = ?";
    return Label::fetchAll<ILabel>( m_dbConnection, req, m_id );
}

int Media::playCount() const
{
    return m_playCount;
}

void Media::increasePlayCount()
{
    m_playCount++;
    m_changed = true;
}

float Media::progress() const
{
    return m_progress;
}

void Media::setProgress( float progress )
{
    if ( progress == m_progress || progress < 0 || progress > 1.0 )
        return;
    m_progress = progress;
    m_changed = true;
}

int Media::rating() const
{
    return m_rating;
}

void Media::setRating( int rating )
{
    if ( m_rating == rating )
        return;
    m_rating = rating;
    m_changed = true;
}

const std::vector<FilePtr>& Media::files() const
{
    auto lock = m_files.lock();
    if ( m_files.isCached() == false )
    {
        static const std::string req = "SELECT * FROM " + policy::FileTable::Name
                + " WHERE media_id = ?";
        m_files = File::fetchAll<IFile>( m_dbConnection, req, m_id );
    }
    return m_files;
}

MoviePtr Media::movie() const
{
    if ( m_subType != SubType::Movie )
        return nullptr;

    auto lock = m_movie.lock();

    if ( m_movie.isCached() == false )
        m_movie = Movie::fromMedia( m_dbConnection, m_id );
    return m_movie.get();
}

void Media::setMovie(MoviePtr movie)
{
    auto lock = m_movie.lock();
    m_movie = movie;
    m_subType = SubType::Movie;
    m_changed = true;
}

bool Media::addVideoTrack(const std::string& codec, unsigned int width, unsigned int height, float fps)
{
    return VideoTrack::create( m_dbConnection, codec, width, height, fps, m_id ) != nullptr;
}

std::vector<VideoTrackPtr> Media::videoTracks()
{
    static const std::string req = "SELECT * FROM " + policy::VideoTrackTable::Name +
            " WHERE media_id = ?";
    return VideoTrack::fetchAll<IVideoTrack>( m_dbConnection, req, m_id );
}

bool Media::addAudioTrack( const std::string& codec, unsigned int bitrate,
                          unsigned int sampleRate, unsigned int nbChannels,
                          const std::string& language, const std::string& desc )
{
    return AudioTrack::create( m_dbConnection, codec, bitrate, sampleRate, nbChannels, language, desc, m_id ) != nullptr;
}

std::vector<AudioTrackPtr> Media::audioTracks()
{
    static const std::string req = "SELECT * FROM " + policy::AudioTrackTable::Name +
            " WHERE media_id = ?";
    return AudioTrack::fetchAll<IAudioTrack>( m_dbConnection, req, m_id );
}

const std::string &Media::thumbnail()
{
    return m_thumbnail;
}

unsigned int Media::insertionDate() const
{
    return m_insertionDate;
}

void Media::setThumbnail(const std::string& thumbnail )
{
    if ( m_thumbnail == thumbnail )
        return;
    m_thumbnail = thumbnail;
    m_changed = true;
}

bool Media::save()
{
    static const std::string req = "UPDATE " + policy::MediaTable::Name + " SET "
            "type = ?, subtype = ?, duration = ?, play_count = ?, progress = ?, rating = ?,"
            "thumbnail = ?, title = ? WHERE id_media = ?";
    if ( m_changed == false )
        return true;
    if ( sqlite::Tools::executeUpdate( m_dbConnection, req, m_type, m_subType, m_duration, m_playCount,
                                       m_progress, m_rating, m_thumbnail, m_title, m_id ) == false )
    {
        return false;
    }
    m_changed = false;
    return true;
}

std::shared_ptr<File> Media::addFile( const fs::IFile& fileFs, Folder& parentFolder, fs::IDirectory& parentFolderFs, IFile::Type type )
{
    auto file = File::create( m_dbConnection, m_id, type, fileFs, parentFolder.id(), parentFolderFs.device()->isRemovable() );
    if ( file == nullptr )
        return nullptr;
    auto lock = m_files.lock();
    if ( m_files.isCached() )
        m_files.get().push_back( file );
    return file;
}

void Media::removeFile( File& file )
{
    file.destroy();
    auto lock = m_files.lock();
    if ( m_files.isCached() == false )
        return;
    m_files.get().erase( std::remove_if( begin( m_files.get() ), end( m_files.get() ), [&file]( const FilePtr& f ) {
        return f->id() == file.id();
    }));
}

unsigned int Media::id() const
{
    return m_id;
}

IMedia::Type Media::type()
{
    return m_type;
}

IMedia::SubType Media::subType() const
{
    return m_subType;
}

void Media::setType( Type type )
{
    if ( m_type != type )
        return;
    m_type = type;
    m_changed = true;
}

const std::string &Media::title() const
{
    return m_title;
}

void Media::setTitle( const std::string &title )
{
    if ( m_title == title )
        return;
    m_title = title;
    m_changed = true;
}

bool Media::createTable( DBConnection connection )
{
    std::string req = "CREATE TABLE IF NOT EXISTS " + policy::MediaTable::Name + "("
            "id_media INTEGER PRIMARY KEY AUTOINCREMENT,"
            "type INTEGER,"
            "subtype INTEGER,"
            "duration INTEGER DEFAULT -1,"
            "play_count UNSIGNED INTEGER,"
            "progress REAL,"
            "rating INTEGER DEFAULT -1,"
            "insertion_date UNSIGNED INTEGER,"
            "thumbnail TEXT,"
            "title TEXT,"
            "is_present BOOLEAN NOT NULL DEFAULT 1"
            ")";
    static const std::string vtableReq = "CREATE VIRTUAL TABLE IF NOT EXISTS "
                + policy::MediaTable::Name + "Fts USING FTS3("
                "title,"
                "labels"
            ")";
    return sqlite::Tools::executeRequest( connection, req ) &&
            sqlite::Tools::executeRequest( connection, vtableReq );
}

bool Media::createTriggers( DBConnection connection )
{
    static const std::string triggerReq = "CREATE TRIGGER IF NOT EXISTS has_files_present AFTER UPDATE OF "
            "is_present ON " + policy::FileTable::Name +
            " BEGIN "
            " UPDATE " + policy::MediaTable::Name + " SET is_present="
                "(SELECT COUNT(id_file) FROM " + policy::FileTable::Name + " WHERE media_id=new.media_id AND is_present=1) "
                "WHERE id_media=new.media_id;"
            " END;";
    static const std::string triggerReq2 = "CREATE TRIGGER IF NOT EXISTS cascade_file_deletion AFTER DELETE ON "
            + policy::FileTable::Name +
            " BEGIN "
            " DELETE FROM " + policy::MediaTable::Name + " WHERE "
                "(SELECT COUNT(id_file) FROM " + policy::FileTable::Name + " WHERE media_id=old.media_id) = 0"
                " AND id_media=old.media_id;"
            " END;";

    static const std::string vtableInsertTrigger = "CREATE TRIGGER IF NOT EXISTS insert_media_fts"
            " AFTER INSERT ON " + policy::MediaTable::Name +
            " BEGIN"
            " INSERT INTO " + policy::MediaTable::Name + "Fts(rowid,title,labels) VALUES(new.id_media, new.title, '');"
            " END";
    static const std::string vtableDeleteTrigger = "CREATE TRIGGER IF NOT EXISTS delete_media_fts"
            " BEFORE DELETE ON " + policy::MediaTable::Name +
            " BEGIN"
            " DELETE FROM " + policy::MediaTable::Name + "Fts WHERE rowid = old.id_media;"
            " END";
    static const std::string vtableUpdateTitleTrigger2 = "CREATE TRIGGER IF NOT EXISTS update_media_title_fts"
              " AFTER UPDATE OF title ON " + policy::MediaTable::Name +
              " BEGIN"
              " UPDATE " + policy::MediaTable::Name + "Fts SET title = new.title WHERE rowid = new.id_media;"
              " END";
    return sqlite::Tools::executeRequest( connection, triggerReq ) &&
            sqlite::Tools::executeRequest( connection, triggerReq2 ) &&
            sqlite::Tools::executeRequest( connection, vtableInsertTrigger ) &&
            sqlite::Tools::executeRequest( connection, vtableDeleteTrigger ) &&
            sqlite::Tools::executeRequest( connection, vtableUpdateTitleTrigger2 );
}

bool Media::addLabel( LabelPtr label )
{
    if ( m_id == 0 || label->id() == 0 )
    {
        LOG_ERROR( "Both file & label need to be inserted in database before being linked together" );
        return false;
    }
    const char* req = "INSERT INTO LabelFileRelation VALUES(?, ?)";
    if ( sqlite::Tools::insert( m_dbConnection, req, label->id(), m_id ) == 0 )
        return false;
    const std::string reqFts = "UPDATE " + policy::MediaTable::Name + "Fts "
        "SET labels = labels || ' ' || ? WHERE rowid = ?";
    return sqlite::Tools::executeUpdate( m_dbConnection, reqFts, label->name(), m_id );
}

bool Media::removeLabel( LabelPtr label )
{
    if ( m_id == 0 || label->id() == 0 )
    {
        LOG_ERROR( "Can't unlink a label/file not inserted in database" );
        return false;
    }
    const char* req = "DELETE FROM LabelFileRelation WHERE label_id = ? AND media_id = ?";
    if ( sqlite::Tools::executeDelete( m_dbConnection, req, label->id(), m_id ) == false )
        return false;
    const std::string reqFts = "UPDATE " + policy::MediaTable::Name + "Fts "
            "SET labels = TRIM(REPLACE(labels, ?, '')) WHERE rowid = ?";
    return sqlite::Tools::executeUpdate( m_dbConnection, reqFts, label->name(), m_id );
}

std::vector<MediaPtr> Media::search( DBConnection dbConn, const std::string& title )
{
    static const std::string req = "SELECT * FROM " + policy::MediaTable::Name + " WHERE"
            " id_media IN (SELECT rowid FROM " + policy::MediaTable::Name + "Fts"
            " WHERE " + policy::MediaTable::Name + "Fts MATCH ?)";
    return Media::fetchAll<IMedia>( dbConn, req, title + "*" );
}
<commit_msg>Media: Don't include non-present media in search results<commit_after>/*****************************************************************************
 * Media Library
 *****************************************************************************
 * Copyright (C) 2015 Hugo Beauzée-Luyssen, Videolabs
 *
 * Authors: Hugo Beauzée-Luyssen<hugo@beauzee.fr>
 *
 * This program 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.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
 * GNU Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public License
 * along with this program; if not, write to the Free Software Foundation,
 * Inc., 51 Franklin Street, Fifth Floor, Boston MA 02110-1301, USA.
 *****************************************************************************/

#include <algorithm>
#include <cassert>
#include <cstdlib>
#include <cstring>
#include <ctime>

#include "Album.h"
#include "AlbumTrack.h"
#include "Artist.h"
#include "AudioTrack.h"
#include "Media.h"
#include "File.h"
#include "Folder.h"
#include "Label.h"
#include "logging/Logger.h"
#include "Movie.h"
#include "ShowEpisode.h"
#include "database/SqliteTools.h"
#include "VideoTrack.h"
#include "filesystem/IFile.h"
#include "filesystem/IDirectory.h"
#include "filesystem/IDevice.h"
#include "utils/Filename.h"

const std::string policy::MediaTable::Name = "Media";
const std::string policy::MediaTable::PrimaryKeyColumn = "id_media";
unsigned int Media::* const policy::MediaTable::PrimaryKey = &Media::m_id;

Media::Media( DBConnection dbConnection, sqlite::Row& row )
    : m_dbConnection( dbConnection )
    , m_changed( false )
{
    row >> m_id
        >> m_type
        >> m_subType
        >> m_duration
        >> m_playCount
        >> m_progress
        >> m_rating
        >> m_insertionDate
        >> m_thumbnail
        >> m_title;
}

Media::Media( const std::string& title, Type type )
    : m_id( 0 )
    , m_type( type )
    , m_subType( SubType::Unknown )
    , m_duration( -1 )
    , m_playCount( 0 )
    , m_progress( .0f )
    , m_rating( -1 )
    , m_insertionDate( time( nullptr ) )
    , m_title( title )
    , m_changed( false )
{
}

std::shared_ptr<Media> Media::create( DBConnection dbConnection, Type type, const fs::IFile* file )
{
    auto self = std::make_shared<Media>( file->name(), type );
    static const std::string req = "INSERT INTO " + policy::MediaTable::Name +
            "(type, insertion_date, title) VALUES(?, ?, ?)";

    if ( insert( dbConnection, self, req, type, self->m_insertionDate, self->m_title ) == false )
        return nullptr;
    self->m_dbConnection = dbConnection;
    return self;
}

AlbumTrackPtr Media::albumTrack() const
{
    if ( m_subType != SubType::AlbumTrack )
        return nullptr;
    auto lock = m_albumTrack.lock();

    if ( m_albumTrack.isCached() == false )
        m_albumTrack = AlbumTrack::fromMedia( m_dbConnection, m_id );
    return m_albumTrack.get();
}

void Media::setAlbumTrack( AlbumTrackPtr albumTrack )
{
    auto lock = m_albumTrack.lock();
    m_albumTrack = albumTrack;
    m_subType = SubType::AlbumTrack;
    m_changed = true;
}

int64_t Media::duration() const
{
    return m_duration;
}

void Media::setDuration( int64_t duration )
{
    if ( m_duration == duration )
        return;
    m_duration = duration;
    m_changed = true;
}

ShowEpisodePtr Media::showEpisode() const
{
    if ( m_subType != SubType::ShowEpisode )
        return nullptr;

    auto lock = m_showEpisode.lock();
    if ( m_showEpisode.isCached() == false )
        m_showEpisode = ShowEpisode::fromMedia( m_dbConnection, m_id );
    return m_showEpisode.get();
}

void Media::setShowEpisode( ShowEpisodePtr episode )
{
    auto lock = m_showEpisode.lock();
    m_showEpisode = episode;
    m_subType = SubType::ShowEpisode;
    m_changed = true;
}

std::vector<LabelPtr> Media::labels()
{
    static const std::string req = "SELECT l.* FROM " + policy::LabelTable::Name + " l "
            "INNER JOIN LabelFileRelation lfr ON lfr.label_id = l.id_label "
            "WHERE lfr.media_id = ?";
    return Label::fetchAll<ILabel>( m_dbConnection, req, m_id );
}

int Media::playCount() const
{
    return m_playCount;
}

void Media::increasePlayCount()
{
    m_playCount++;
    m_changed = true;
}

float Media::progress() const
{
    return m_progress;
}

void Media::setProgress( float progress )
{
    if ( progress == m_progress || progress < 0 || progress > 1.0 )
        return;
    m_progress = progress;
    m_changed = true;
}

int Media::rating() const
{
    return m_rating;
}

void Media::setRating( int rating )
{
    if ( m_rating == rating )
        return;
    m_rating = rating;
    m_changed = true;
}

const std::vector<FilePtr>& Media::files() const
{
    auto lock = m_files.lock();
    if ( m_files.isCached() == false )
    {
        static const std::string req = "SELECT * FROM " + policy::FileTable::Name
                + " WHERE media_id = ?";
        m_files = File::fetchAll<IFile>( m_dbConnection, req, m_id );
    }
    return m_files;
}

MoviePtr Media::movie() const
{
    if ( m_subType != SubType::Movie )
        return nullptr;

    auto lock = m_movie.lock();

    if ( m_movie.isCached() == false )
        m_movie = Movie::fromMedia( m_dbConnection, m_id );
    return m_movie.get();
}

void Media::setMovie(MoviePtr movie)
{
    auto lock = m_movie.lock();
    m_movie = movie;
    m_subType = SubType::Movie;
    m_changed = true;
}

bool Media::addVideoTrack(const std::string& codec, unsigned int width, unsigned int height, float fps)
{
    return VideoTrack::create( m_dbConnection, codec, width, height, fps, m_id ) != nullptr;
}

std::vector<VideoTrackPtr> Media::videoTracks()
{
    static const std::string req = "SELECT * FROM " + policy::VideoTrackTable::Name +
            " WHERE media_id = ?";
    return VideoTrack::fetchAll<IVideoTrack>( m_dbConnection, req, m_id );
}

bool Media::addAudioTrack( const std::string& codec, unsigned int bitrate,
                          unsigned int sampleRate, unsigned int nbChannels,
                          const std::string& language, const std::string& desc )
{
    return AudioTrack::create( m_dbConnection, codec, bitrate, sampleRate, nbChannels, language, desc, m_id ) != nullptr;
}

std::vector<AudioTrackPtr> Media::audioTracks()
{
    static const std::string req = "SELECT * FROM " + policy::AudioTrackTable::Name +
            " WHERE media_id = ?";
    return AudioTrack::fetchAll<IAudioTrack>( m_dbConnection, req, m_id );
}

const std::string &Media::thumbnail()
{
    return m_thumbnail;
}

unsigned int Media::insertionDate() const
{
    return m_insertionDate;
}

void Media::setThumbnail(const std::string& thumbnail )
{
    if ( m_thumbnail == thumbnail )
        return;
    m_thumbnail = thumbnail;
    m_changed = true;
}

bool Media::save()
{
    static const std::string req = "UPDATE " + policy::MediaTable::Name + " SET "
            "type = ?, subtype = ?, duration = ?, play_count = ?, progress = ?, rating = ?,"
            "thumbnail = ?, title = ? WHERE id_media = ?";
    if ( m_changed == false )
        return true;
    if ( sqlite::Tools::executeUpdate( m_dbConnection, req, m_type, m_subType, m_duration, m_playCount,
                                       m_progress, m_rating, m_thumbnail, m_title, m_id ) == false )
    {
        return false;
    }
    m_changed = false;
    return true;
}

std::shared_ptr<File> Media::addFile( const fs::IFile& fileFs, Folder& parentFolder, fs::IDirectory& parentFolderFs, IFile::Type type )
{
    auto file = File::create( m_dbConnection, m_id, type, fileFs, parentFolder.id(), parentFolderFs.device()->isRemovable() );
    if ( file == nullptr )
        return nullptr;
    auto lock = m_files.lock();
    if ( m_files.isCached() )
        m_files.get().push_back( file );
    return file;
}

void Media::removeFile( File& file )
{
    file.destroy();
    auto lock = m_files.lock();
    if ( m_files.isCached() == false )
        return;
    m_files.get().erase( std::remove_if( begin( m_files.get() ), end( m_files.get() ), [&file]( const FilePtr& f ) {
        return f->id() == file.id();
    }));
}

unsigned int Media::id() const
{
    return m_id;
}

IMedia::Type Media::type()
{
    return m_type;
}

IMedia::SubType Media::subType() const
{
    return m_subType;
}

void Media::setType( Type type )
{
    if ( m_type != type )
        return;
    m_type = type;
    m_changed = true;
}

const std::string &Media::title() const
{
    return m_title;
}

void Media::setTitle( const std::string &title )
{
    if ( m_title == title )
        return;
    m_title = title;
    m_changed = true;
}

bool Media::createTable( DBConnection connection )
{
    std::string req = "CREATE TABLE IF NOT EXISTS " + policy::MediaTable::Name + "("
            "id_media INTEGER PRIMARY KEY AUTOINCREMENT,"
            "type INTEGER,"
            "subtype INTEGER,"
            "duration INTEGER DEFAULT -1,"
            "play_count UNSIGNED INTEGER,"
            "progress REAL,"
            "rating INTEGER DEFAULT -1,"
            "insertion_date UNSIGNED INTEGER,"
            "thumbnail TEXT,"
            "title TEXT,"
            "is_present BOOLEAN NOT NULL DEFAULT 1"
            ")";
    static const std::string vtableReq = "CREATE VIRTUAL TABLE IF NOT EXISTS "
                + policy::MediaTable::Name + "Fts USING FTS3("
                "title,"
                "labels"
            ")";
    return sqlite::Tools::executeRequest( connection, req ) &&
            sqlite::Tools::executeRequest( connection, vtableReq );
}

bool Media::createTriggers( DBConnection connection )
{
    static const std::string triggerReq = "CREATE TRIGGER IF NOT EXISTS has_files_present AFTER UPDATE OF "
            "is_present ON " + policy::FileTable::Name +
            " BEGIN "
            " UPDATE " + policy::MediaTable::Name + " SET is_present="
                "(SELECT COUNT(id_file) FROM " + policy::FileTable::Name + " WHERE media_id=new.media_id AND is_present=1) "
                "WHERE id_media=new.media_id;"
            " END;";
    static const std::string triggerReq2 = "CREATE TRIGGER IF NOT EXISTS cascade_file_deletion AFTER DELETE ON "
            + policy::FileTable::Name +
            " BEGIN "
            " DELETE FROM " + policy::MediaTable::Name + " WHERE "
                "(SELECT COUNT(id_file) FROM " + policy::FileTable::Name + " WHERE media_id=old.media_id) = 0"
                " AND id_media=old.media_id;"
            " END;";

    static const std::string vtableInsertTrigger = "CREATE TRIGGER IF NOT EXISTS insert_media_fts"
            " AFTER INSERT ON " + policy::MediaTable::Name +
            " BEGIN"
            " INSERT INTO " + policy::MediaTable::Name + "Fts(rowid,title,labels) VALUES(new.id_media, new.title, '');"
            " END";
    static const std::string vtableDeleteTrigger = "CREATE TRIGGER IF NOT EXISTS delete_media_fts"
            " BEFORE DELETE ON " + policy::MediaTable::Name +
            " BEGIN"
            " DELETE FROM " + policy::MediaTable::Name + "Fts WHERE rowid = old.id_media;"
            " END";
    static const std::string vtableUpdateTitleTrigger2 = "CREATE TRIGGER IF NOT EXISTS update_media_title_fts"
              " AFTER UPDATE OF title ON " + policy::MediaTable::Name +
              " BEGIN"
              " UPDATE " + policy::MediaTable::Name + "Fts SET title = new.title WHERE rowid = new.id_media;"
              " END";
    return sqlite::Tools::executeRequest( connection, triggerReq ) &&
            sqlite::Tools::executeRequest( connection, triggerReq2 ) &&
            sqlite::Tools::executeRequest( connection, vtableInsertTrigger ) &&
            sqlite::Tools::executeRequest( connection, vtableDeleteTrigger ) &&
            sqlite::Tools::executeRequest( connection, vtableUpdateTitleTrigger2 );
}

bool Media::addLabel( LabelPtr label )
{
    if ( m_id == 0 || label->id() == 0 )
    {
        LOG_ERROR( "Both file & label need to be inserted in database before being linked together" );
        return false;
    }
    const char* req = "INSERT INTO LabelFileRelation VALUES(?, ?)";
    if ( sqlite::Tools::insert( m_dbConnection, req, label->id(), m_id ) == 0 )
        return false;
    const std::string reqFts = "UPDATE " + policy::MediaTable::Name + "Fts "
        "SET labels = labels || ' ' || ? WHERE rowid = ?";
    return sqlite::Tools::executeUpdate( m_dbConnection, reqFts, label->name(), m_id );
}

bool Media::removeLabel( LabelPtr label )
{
    if ( m_id == 0 || label->id() == 0 )
    {
        LOG_ERROR( "Can't unlink a label/file not inserted in database" );
        return false;
    }
    const char* req = "DELETE FROM LabelFileRelation WHERE label_id = ? AND media_id = ?";
    if ( sqlite::Tools::executeDelete( m_dbConnection, req, label->id(), m_id ) == false )
        return false;
    const std::string reqFts = "UPDATE " + policy::MediaTable::Name + "Fts "
            "SET labels = TRIM(REPLACE(labels, ?, '')) WHERE rowid = ?";
    return sqlite::Tools::executeUpdate( m_dbConnection, reqFts, label->name(), m_id );
}

std::vector<MediaPtr> Media::search( DBConnection dbConn, const std::string& title )
{
    static const std::string req = "SELECT * FROM " + policy::MediaTable::Name + " WHERE"
            " id_media IN (SELECT rowid FROM " + policy::MediaTable::Name + "Fts"
            " WHERE " + policy::MediaTable::Name + "Fts MATCH ?)"
            "AND is_present = 1";
    return Media::fetchAll<IMedia>( dbConn, req, title + "*" );
}
<|endoftext|>


<commit_before>/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.

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

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

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

#include <string>
#include <utility>

#include "tensorflow/core/data/name_utils.h"
#include "tensorflow/core/data/split_utils.h"
#include "tensorflow/core/framework/partial_tensor_shape.h"
#include "tensorflow/core/framework/tensor.h"

namespace tensorflow {
namespace data {

// See documentation in ../../ops/dataset_ops.cc for a high-level
// description of the following op.

/* static */ constexpr const char* const ConcatenateDatasetOp::kDatasetType;
/* static */ constexpr const char* const ConcatenateDatasetOp::kInputDataset;
/* static */ constexpr const char* const ConcatenateDatasetOp::kAnotherDataset;
/* static */ constexpr const char* const ConcatenateDatasetOp::kOutputTypes;
/* static */ constexpr const char* const ConcatenateDatasetOp::kOutputShapes;

constexpr char kIndex[] = "i";
constexpr char kInputImplUninitialized[] = "input_impl_uninitialized";

class ConcatenateDatasetOp::Dataset : public DatasetBase {
 public:
  explicit Dataset(OpKernelContext* ctx, const DatasetBase* input,
                   const DatasetBase* to_concatenate)
      : DatasetBase(DatasetContext(ctx)),
        input_(input),
        to_concatenate_(to_concatenate),
        input_cardinality_(input->Cardinality()),
        to_concatenate_cardinality_(to_concatenate_->Cardinality()) {
    input_->Ref();
    to_concatenate_->Ref();

    auto os_input = input->output_shapes();
    auto os_concatenate = to_concatenate->output_shapes();
    PartialTensorShape output_tensorshape({});
    for (int i = 0; i < os_input.size(); i++) {
      OP_REQUIRES_OK(ctx, MostSpecificCompatibleShape(os_input[i], os_concatenate[i],
                                        &output_tensorshape));
      output_shapes_.push_back(output_tensorshape);
    }
  }
  ~Dataset() override {
    input_->Unref();
    to_concatenate_->Unref();
  }

  std::unique_ptr<IteratorBase> MakeIteratorInternal(
      const string& prefix) const override {
    return std::make_unique<Iterator>(Iterator::Params{
        this, name_utils::IteratorPrefix(kDatasetType, prefix)});
  }

  Status MakeSplitProviders(std::vector<std::unique_ptr<SplitProvider>>*
                                split_providers) const override {
    TF_ASSIGN_OR_RETURN(*split_providers, GetSplitProviders(this));
    return OkStatus();
  }

  const DataTypeVector& output_dtypes() const override {
    return input_->output_dtypes();
  }

  const std::vector<PartialTensorShape>& output_shapes() const override {
    return output_shapes_;
  }

  string DebugString() const override {
    return name_utils::DatasetDebugString(kDatasetType);
  }

  int64_t CardinalityInternal() const override {
    if (input_cardinality_ == kInfiniteCardinality ||
        to_concatenate_cardinality_ == kInfiniteCardinality) {
      return kInfiniteCardinality;
    }
    if (input_cardinality_ == kUnknownCardinality ||
        to_concatenate_cardinality_ == kUnknownCardinality) {
      return kUnknownCardinality;
    }
    return input_cardinality_ + to_concatenate_cardinality_;
  }

  int64_t CardinalityInternal(CardinalityOptions options) const override {
    int64_t input_cardinality = input_->Cardinality(options);
    int64_t to_concatenate_cardinality = to_concatenate_->Cardinality(options);

    if (input_cardinality == kInfiniteCardinality ||
        to_concatenate_cardinality == kInfiniteCardinality) {
      return kInfiniteCardinality;
    }
    if (input_cardinality == kUnknownCardinality ||
        to_concatenate_cardinality == kUnknownCardinality) {
      return kUnknownCardinality;
    }
    return input_cardinality + to_concatenate_cardinality;
  }

  Status InputDatasets(std::vector<const DatasetBase*>* inputs) const override {
    inputs->push_back(input_);
    inputs->push_back(to_concatenate_);
    return OkStatus();
  }

  Status CheckExternalState() const override {
    TF_RETURN_IF_ERROR(input_->CheckExternalState());
    return to_concatenate_->CheckExternalState();
  }

  Status Get(OpKernelContext* ctx, int64 index,
             std::vector<Tensor>* out_tensors) const override {
    TF_RETURN_IF_ERROR(CheckRandomAccessCompatible(index));
    if (index < input_cardinality_) {
      TF_RETURN_IF_ERROR(input_->Get(ctx, index, out_tensors));
    } else {
      TF_RETURN_IF_ERROR(
          to_concatenate_->Get(ctx, index - input_cardinality_, out_tensors));
    }
    return OkStatus();
  }

 protected:
  Status AsGraphDefInternal(SerializationContext* ctx,
                            DatasetGraphDefBuilder* b,
                            Node** output) const override {
    Node* input_graph = nullptr;
    TF_RETURN_IF_ERROR(b->AddInputDataset(ctx, input_, &input_graph));
    Node* to_concatenate_graph = nullptr;
    TF_RETURN_IF_ERROR(
        b->AddInputDataset(ctx, to_concatenate_, &to_concatenate_graph));
    TF_RETURN_IF_ERROR(
        b->AddDataset(this, {input_graph, to_concatenate_graph}, output));
    return OkStatus();
  }

 private:
  class Iterator : public DatasetIterator<Dataset> {
   public:
    explicit Iterator(const Params& params)
        : DatasetIterator<Dataset>(params), i_(0) {}

    Status Initialize(IteratorContext* ctx) override {
      TF_ASSIGN_OR_RETURN(input_contexts_,
                          CreateInputIteratorContexts(ctx, dataset()));
      return dataset()->input_->MakeIterator(&input_contexts_[0], this,
                                             strings::StrCat(prefix(), "[0]"),
                                             &input_impl_);
    }

    Status GetNextInternal(IteratorContext* ctx,
                           std::vector<Tensor>* out_tensors,
                           bool* end_of_sequence) override {
      mutex_lock l(mu_);
      if (!input_impl_) {
        *end_of_sequence = true;
        return OkStatus();
      }
      while (i_ < 2) {
        TF_RETURN_IF_ERROR(input_impl_->GetNext(&input_contexts_[i_],
                                                out_tensors, end_of_sequence));
        if (!*end_of_sequence) {
          return OkStatus();
        }
        if (++i_ < 2) {
          TF_RETURN_IF_ERROR(dataset()->to_concatenate_->MakeIterator(
              &input_contexts_[i_], this, strings::StrCat(prefix(), "[1]"),
              &input_impl_));
        }
      }
      *end_of_sequence = true;
      input_impl_.reset();
      return OkStatus();
    }

   protected:
    std::shared_ptr<model::Node> CreateNode(
        IteratorContext* ctx, model::Node::Args args) const override {
      return model::MakeKnownRatioNode(std::move(args),
                                       /*ratio=*/1);
    }

    Status SaveInternal(SerializationContext* ctx,
                        IteratorStateWriter* writer) override {
      mutex_lock l(mu_);
      TF_RETURN_IF_ERROR(writer->WriteScalar(full_name(kIndex), i_));
      if (input_impl_) {
        TF_RETURN_IF_ERROR(SaveInput(ctx, writer, input_impl_));
      } else {
        TF_RETURN_IF_ERROR(
            writer->WriteScalar(full_name(kInputImplUninitialized), ""));
      }
      return OkStatus();
    }

    Status RestoreInternal(IteratorContext* ctx,
                           IteratorStateReader* reader) override {
      mutex_lock l(mu_);
      TF_RETURN_IF_ERROR(reader->ReadScalar(full_name(kIndex), &i_));
      if (reader->Contains(full_name(kInputImplUninitialized))) {
        input_impl_.reset();
        return OkStatus();
      }
      if (!TF_PREDICT_TRUE(i_ >= 0 && i_ <= 2))
        return errors::InvalidArgument("i_ must be in range [0, 2].");
      if (i_ == 1) {
        TF_RETURN_IF_ERROR(dataset()->to_concatenate_->MakeIterator(
            ctx, this, strings::StrCat(prefix(), "[1]"), &input_impl_));
      } else if (i_ == 2) {
        input_impl_.reset();
      }
      if (input_impl_) {
        TF_RETURN_IF_ERROR(RestoreInput(ctx, reader, input_impl_));
      }
      return OkStatus();
    }

   private:
    mutex mu_;
    int64_t i_ TF_GUARDED_BY(mu_);
    std::unique_ptr<IteratorBase> input_impl_ TF_GUARDED_BY(mu_);
    std::vector<IteratorContext> input_contexts_;
  };

  Status MostSpecificCompatibleShape(
      const PartialTensorShape& ts1, const PartialTensorShape& ts2,
      std::unique_ptr<PartialTensorShape>* output_tensorshape) {
    if (ts1.dims() != ts2.dims() || ts1.unknown_rank() || ts2.unknown_rank())
      *output_tensorshape = std::make_unique<PartialTensorShape>();
      return OkStatus();
    *output_tensorshape =
        std::make_unique<PartialTensorShape>( std::initializer_list<int64_t>({}));
    auto dims1 = ts1.dim_sizes();
    auto dims2 = ts2.dim_sizes();
    for (int d = 0; d < ts1.dims(); d++) {
      if (dims1[d] == dims2[d])
        TF_RETURN_IF_ERROR(
            (*output_tensorshape)->AddDimWithStatus(dims1[d]));
      else
        (*output_tensorshape)->AddDim(-1);
    }
    return OkStatus();
  }

  const DatasetBase* input_;
  const DatasetBase* to_concatenate_;
  const int64_t input_cardinality_;
  const int64_t to_concatenate_cardinality_;
  std::vector<PartialTensorShape> output_shapes_;
};

ConcatenateDatasetOp::ConcatenateDatasetOp(OpKernelConstruction* ctx)
    : BinaryDatasetOpKernel(ctx) {}

void ConcatenateDatasetOp::MakeDataset(OpKernelContext* ctx, DatasetBase* input,
                                       DatasetBase* to_concatenate,
                                       DatasetBase** output) {
  OP_REQUIRES(ctx, input->output_dtypes() == to_concatenate->output_dtypes(),
              errors::InvalidArgument(
                  "input dataset and dataset to concatenate"
                  " have different output_types %s and %s",
                  (DataTypeVectorString(input->output_dtypes()),
                   DataTypeVectorString(to_concatenate->output_dtypes()))));
  *output = new Dataset(ctx, input, to_concatenate);
}

namespace {
REGISTER_KERNEL_BUILDER(Name("ConcatenateDataset").Device(DEVICE_CPU),
                        ConcatenateDatasetOp);
}  // namespace
}  // namespace data
}  // namespace tensorflow
<commit_msg>Update tensorflow/core/kernels/data/concatenate_dataset_op.cc<commit_after>/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.

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

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

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

#include <string>
#include <utility>

#include "tensorflow/core/data/name_utils.h"
#include "tensorflow/core/data/split_utils.h"
#include "tensorflow/core/framework/partial_tensor_shape.h"
#include "tensorflow/core/framework/tensor.h"

namespace tensorflow {
namespace data {

// See documentation in ../../ops/dataset_ops.cc for a high-level
// description of the following op.

/* static */ constexpr const char* const ConcatenateDatasetOp::kDatasetType;
/* static */ constexpr const char* const ConcatenateDatasetOp::kInputDataset;
/* static */ constexpr const char* const ConcatenateDatasetOp::kAnotherDataset;
/* static */ constexpr const char* const ConcatenateDatasetOp::kOutputTypes;
/* static */ constexpr const char* const ConcatenateDatasetOp::kOutputShapes;

constexpr char kIndex[] = "i";
constexpr char kInputImplUninitialized[] = "input_impl_uninitialized";

class ConcatenateDatasetOp::Dataset : public DatasetBase {
 public:
  explicit Dataset(OpKernelContext* ctx, const DatasetBase* input,
                   const DatasetBase* to_concatenate)
      : DatasetBase(DatasetContext(ctx)),
        input_(input),
        to_concatenate_(to_concatenate),
        input_cardinality_(input->Cardinality()),
        to_concatenate_cardinality_(to_concatenate_->Cardinality()) {
    input_->Ref();
    to_concatenate_->Ref();

    auto os_input = input->output_shapes();
    auto os_concatenate = to_concatenate->output_shapes();
    PartialTensorShape output_tensorshape({});
    for (int i = 0; i < os_input.size(); i++) {
      OP_REQUIRES_OK(ctx, MostSpecificCompatibleShape(os_input[i], os_concatenate[i],
                                        &output_tensorshape));
      output_shapes_.push_back(output_tensorshape);
    }
  }
  ~Dataset() override {
    input_->Unref();
    to_concatenate_->Unref();
  }

  std::unique_ptr<IteratorBase> MakeIteratorInternal(
      const string& prefix) const override {
    return std::make_unique<Iterator>(Iterator::Params{
        this, name_utils::IteratorPrefix(kDatasetType, prefix)});
  }

  Status MakeSplitProviders(std::vector<std::unique_ptr<SplitProvider>>*
                                split_providers) const override {
    TF_ASSIGN_OR_RETURN(*split_providers, GetSplitProviders(this));
    return OkStatus();
  }

  const DataTypeVector& output_dtypes() const override {
    return input_->output_dtypes();
  }

  const std::vector<PartialTensorShape>& output_shapes() const override {
    return output_shapes_;
  }

  string DebugString() const override {
    return name_utils::DatasetDebugString(kDatasetType);
  }

  int64_t CardinalityInternal() const override {
    if (input_cardinality_ == kInfiniteCardinality ||
        to_concatenate_cardinality_ == kInfiniteCardinality) {
      return kInfiniteCardinality;
    }
    if (input_cardinality_ == kUnknownCardinality ||
        to_concatenate_cardinality_ == kUnknownCardinality) {
      return kUnknownCardinality;
    }
    return input_cardinality_ + to_concatenate_cardinality_;
  }

  int64_t CardinalityInternal(CardinalityOptions options) const override {
    int64_t input_cardinality = input_->Cardinality(options);
    int64_t to_concatenate_cardinality = to_concatenate_->Cardinality(options);

    if (input_cardinality == kInfiniteCardinality ||
        to_concatenate_cardinality == kInfiniteCardinality) {
      return kInfiniteCardinality;
    }
    if (input_cardinality == kUnknownCardinality ||
        to_concatenate_cardinality == kUnknownCardinality) {
      return kUnknownCardinality;
    }
    return input_cardinality + to_concatenate_cardinality;
  }

  Status InputDatasets(std::vector<const DatasetBase*>* inputs) const override {
    inputs->push_back(input_);
    inputs->push_back(to_concatenate_);
    return OkStatus();
  }

  Status CheckExternalState() const override {
    TF_RETURN_IF_ERROR(input_->CheckExternalState());
    return to_concatenate_->CheckExternalState();
  }

  Status Get(OpKernelContext* ctx, int64 index,
             std::vector<Tensor>* out_tensors) const override {
    TF_RETURN_IF_ERROR(CheckRandomAccessCompatible(index));
    if (index < input_cardinality_) {
      TF_RETURN_IF_ERROR(input_->Get(ctx, index, out_tensors));
    } else {
      TF_RETURN_IF_ERROR(
          to_concatenate_->Get(ctx, index - input_cardinality_, out_tensors));
    }
    return OkStatus();
  }

 protected:
  Status AsGraphDefInternal(SerializationContext* ctx,
                            DatasetGraphDefBuilder* b,
                            Node** output) const override {
    Node* input_graph = nullptr;
    TF_RETURN_IF_ERROR(b->AddInputDataset(ctx, input_, &input_graph));
    Node* to_concatenate_graph = nullptr;
    TF_RETURN_IF_ERROR(
        b->AddInputDataset(ctx, to_concatenate_, &to_concatenate_graph));
    TF_RETURN_IF_ERROR(
        b->AddDataset(this, {input_graph, to_concatenate_graph}, output));
    return OkStatus();
  }

 private:
  class Iterator : public DatasetIterator<Dataset> {
   public:
    explicit Iterator(const Params& params)
        : DatasetIterator<Dataset>(params), i_(0) {}

    Status Initialize(IteratorContext* ctx) override {
      TF_ASSIGN_OR_RETURN(input_contexts_,
                          CreateInputIteratorContexts(ctx, dataset()));
      return dataset()->input_->MakeIterator(&input_contexts_[0], this,
                                             strings::StrCat(prefix(), "[0]"),
                                             &input_impl_);
    }

    Status GetNextInternal(IteratorContext* ctx,
                           std::vector<Tensor>* out_tensors,
                           bool* end_of_sequence) override {
      mutex_lock l(mu_);
      if (!input_impl_) {
        *end_of_sequence = true;
        return OkStatus();
      }
      while (i_ < 2) {
        TF_RETURN_IF_ERROR(input_impl_->GetNext(&input_contexts_[i_],
                                                out_tensors, end_of_sequence));
        if (!*end_of_sequence) {
          return OkStatus();
        }
        if (++i_ < 2) {
          TF_RETURN_IF_ERROR(dataset()->to_concatenate_->MakeIterator(
              &input_contexts_[i_], this, strings::StrCat(prefix(), "[1]"),
              &input_impl_));
        }
      }
      *end_of_sequence = true;
      input_impl_.reset();
      return OkStatus();
    }

   protected:
    std::shared_ptr<model::Node> CreateNode(
        IteratorContext* ctx, model::Node::Args args) const override {
      return model::MakeKnownRatioNode(std::move(args),
                                       /*ratio=*/1);
    }

    Status SaveInternal(SerializationContext* ctx,
                        IteratorStateWriter* writer) override {
      mutex_lock l(mu_);
      TF_RETURN_IF_ERROR(writer->WriteScalar(full_name(kIndex), i_));
      if (input_impl_) {
        TF_RETURN_IF_ERROR(SaveInput(ctx, writer, input_impl_));
      } else {
        TF_RETURN_IF_ERROR(
            writer->WriteScalar(full_name(kInputImplUninitialized), ""));
      }
      return OkStatus();
    }

    Status RestoreInternal(IteratorContext* ctx,
                           IteratorStateReader* reader) override {
      mutex_lock l(mu_);
      TF_RETURN_IF_ERROR(reader->ReadScalar(full_name(kIndex), &i_));
      if (reader->Contains(full_name(kInputImplUninitialized))) {
        input_impl_.reset();
        return OkStatus();
      }
      if (!TF_PREDICT_TRUE(i_ >= 0 && i_ <= 2))
        return errors::InvalidArgument("i_ must be in range [0, 2].");
      if (i_ == 1) {
        TF_RETURN_IF_ERROR(dataset()->to_concatenate_->MakeIterator(
            ctx, this, strings::StrCat(prefix(), "[1]"), &input_impl_));
      } else if (i_ == 2) {
        input_impl_.reset();
      }
      if (input_impl_) {
        TF_RETURN_IF_ERROR(RestoreInput(ctx, reader, input_impl_));
      }
      return OkStatus();
    }

   private:
    mutex mu_;
    int64_t i_ TF_GUARDED_BY(mu_);
    std::unique_ptr<IteratorBase> input_impl_ TF_GUARDED_BY(mu_);
    std::vector<IteratorContext> input_contexts_;
  };

  Status MostSpecificCompatibleShape(
      const PartialTensorShape& ts1, const PartialTensorShape& ts2,
      PartialTensorShape* output_tensorshape) {
    if (ts1.dims() != ts2.dims() || ts1.unknown_rank() || ts2.unknown_rank())
      *output_tensorshape = std::make_unique<PartialTensorShape>();
      return OkStatus();
    *output_tensorshape =
        std::make_unique<PartialTensorShape>( std::initializer_list<int64_t>({}));
    auto dims1 = ts1.dim_sizes();
    auto dims2 = ts2.dim_sizes();
    for (int d = 0; d < ts1.dims(); d++) {
      if (dims1[d] == dims2[d])
        TF_RETURN_IF_ERROR(
            (*output_tensorshape)->AddDimWithStatus(dims1[d]));
      else
        (*output_tensorshape)->AddDim(-1);
    }
    return OkStatus();
  }

  const DatasetBase* input_;
  const DatasetBase* to_concatenate_;
  const int64_t input_cardinality_;
  const int64_t to_concatenate_cardinality_;
  std::vector<PartialTensorShape> output_shapes_;
};

ConcatenateDatasetOp::ConcatenateDatasetOp(OpKernelConstruction* ctx)
    : BinaryDatasetOpKernel(ctx) {}

void ConcatenateDatasetOp::MakeDataset(OpKernelContext* ctx, DatasetBase* input,
                                       DatasetBase* to_concatenate,
                                       DatasetBase** output) {
  OP_REQUIRES(ctx, input->output_dtypes() == to_concatenate->output_dtypes(),
              errors::InvalidArgument(
                  "input dataset and dataset to concatenate"
                  " have different output_types %s and %s",
                  (DataTypeVectorString(input->output_dtypes()),
                   DataTypeVectorString(to_concatenate->output_dtypes()))));
  *output = new Dataset(ctx, input, to_concatenate);
}

namespace {
REGISTER_KERNEL_BUILDER(Name("ConcatenateDataset").Device(DEVICE_CPU),
                        ConcatenateDatasetOp);
}  // namespace
}  // namespace data
}  // namespace tensorflow
<|endoftext|>