<commit_before>/* ***** BEGIN LICENSE BLOCK *****
 * Version: MPL 1.1/GPL 2.0/LGPL 2.1
 *
 * The contents of this file are subject to the Mozilla Public License Version
 * 1.1 (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.mozilla.org/MPL/
 *
 * Software distributed under the License is distributed on an "AS IS" basis,
 * WITHOUT WARRANTY OF ANY KIND, either express or implied. See the License
 * for the specific language governing rights and limitations under the
 * License.
 *
 * The Original Code is [Open Source Virtual Machine.].
 *
 * The Initial Developer of the Original Code is
 * Adobe System Incorporated.
 * Portions created by the Initial Developer are Copyright (C) 2004-2006
 * the Initial Developer. All Rights Reserved.
 *
 * Contributor(s):
 *   Adobe AS3 Team
 *
 * Alternatively, the contents of this file may be used under the terms of
 * either the GNU General Public License Version 2 or later (the "GPL"), or
 * the GNU Lesser General Public License Version 2.1 or later (the "LGPL"),
 * in which case the provisions of the GPL or the LGPL are applicable instead
 * of those above. If you wish to allow use of your version of this file only
 * under the terms of either the GPL or the LGPL, and not to allow others to
 * use your version of this file under the terms of the MPL, indicate your
 * decision by deleting the provisions above and replace them with the notice
 * and other provisions required by the GPL or the LGPL. If you do not delete
 * the provisions above, a recipient may use your version of this file under
 * the terms of any one of the MPL, the GPL or the LGPL.
 *
 * ***** END LICENSE BLOCK ***** */

#ifdef AVMSHELL_BUILD
#include "avmshell.h"
#else
// player
#include "platformbuild.h"
#include "avmplayer.h"
#include "DictionaryGlue.h"
#endif

using namespace MMgc;

namespace avmplus
{
	DictionaryClass::DictionaryClass(VTable *vtable)
	: ClassClosure(vtable)
	{
		createVanillaPrototype();
		vtable->traits->itraits->isDictionary = true;
	}

	ScriptObject *DictionaryClass::createInstance(VTable *ivtable, ScriptObject* /*delegate*/)
	{
 		GCAssert(ivtable->traits->isDictionary == true);
		return new (core()->GetGC(), ivtable->getExtraSize()) DictionaryObject(ivtable, prototype);
	}

	DictionaryObject::DictionaryObject(VTable *vtable, ScriptObject *delegate)
		: ScriptObject(vtable, delegate)
	{
		GCAssert(vtable->traits->isDictionary == true);
	}

	void DictionaryObject::init(bool weakKeys)
	{
		GCAssert(vtable->traits->isDictionary == true);
		MMgc::GC* gc = this->gc();

		HeapHashtable* ht = weakKeys ? new (gc) WeakKeyHashtable(gc)
									: new (gc) HeapHashtable(gc);

		//store pointer of newly created hashtable, encapsulated with writebarrier,
		//at the hashtable offset address of the corresponding traits
		union {
			uint8_t* p;
			HeapHashtable** hht;
		};
		p = (uint8_t*)this + vtable->traits->getHashtableOffset();
		WB(gc, this, hht, ht);
	}

	Atom FASTCALL DictionaryObject::getKeyFromObject(Atom key) const
	{
		AvmAssert(AvmCore::isObject(key));
		ScriptObject* obj = AvmCore::atomToScriptObject(key);
		AvmAssert(Traits::getBuiltinType(obj->traits()) != BUILTIN_qName);
		AvmAssert(MMgc::GC::Size(obj) >= sizeof(ScriptObject));
		(void)obj;

		// FIXME: this doesn't work, need to convert back to an XMLObject
		// on the way out or intern XMLObject's somehow
		//if(AvmCore::isXML(key))
		//	key = AvmCore::genericObjectToAtom(AvmCore::atomToXML(key));

		return key;
	}

	Atom DictionaryObject::getAtomProperty(Atom key) const
	{
		if (AvmCore::isObject(key))
		{
			return getHeapHashtable()->get(getKeyFromObject(key));
		}

		return ScriptObject::getAtomProperty(key);
	}

	bool DictionaryObject::hasAtomProperty(Atom key) const
	{
		if (AvmCore::isObject(key))
		{
			return getHeapHashtable()->contains(getKeyFromObject(key));
		}

		return ScriptObject::hasAtomProperty(key);
	}

	bool DictionaryObject::deleteAtomProperty(Atom key)
	{
		if (AvmCore::isObject(key))
		{
			getHeapHashtable()->remove(getKeyFromObject(key));
			return true;
		}

		return ScriptObject::deleteAtomProperty(key);
	}

	void DictionaryObject::setAtomProperty(Atom key, Atom value)
	{
		if (AvmCore::isObject(key))
		{
			getHeapHashtable()->add(getKeyFromObject(key), value);
			return;
		}

		ScriptObject::setAtomProperty(key, value);
	}

	Atom DictionaryObject::nextName(int index)
	{
		Atom k = ScriptObject::nextName(index);

		if (AvmCore::isGenericObject(k) && getHeapHashtable()->weakKeys())
		{
			GCWeakRef* ref = (GCWeakRef*)AvmCore::atomToGenericObject(k);
			union {
				GCObject* key_o;
				ScriptObject* key;
			};
			key_o = ref->get();
			if (key)
			{
				AvmAssert(key->traits() != NULL);
				return key->atom();
			}

			return undefinedAtom;
		}

		return k;
	}

	int DictionaryObject::nextNameIndex(int index)
	{
		AvmAssert(index >= 0);

		// this can happen if you break in debugger in a subclasses constructor before super
		// has been called -- let's do it in all builds, it's better than crashing.
		HeapHashtable* hht = getHeapHashtable();
		if (!hht)
		{
			return 0;
		}

		// Advance to first non-empty slot.
		index <<= 1;
		InlineHashtable* const ht = hht->get_ht();
		Atom* const atoms = ht->getAtoms();
		int const numAtoms = ht->getCapacity();
		if (hht->weakKeys())
		{
			while (index < numAtoms)
			{
				Atom const a = atoms[index];
				if (AvmCore::isGenericObject(a))
				{
					GCWeakRef *weakRef = (GCWeakRef*)AvmCore::atomToGenericObject(a);
					if (weakRef->get())
						return (index>>1)+1;

					atoms[index] = InlineHashtable::DELETED;
					atoms[index+1] = InlineHashtable::DELETED;
					ht->setHasDeletedItems();
				}
				else if (a != 0 && a != InlineHashtable::DELETED)
				{
					return (index>>1)+1;
				}
				index += 2;
			}
		}
		else
		{
			while (index < numAtoms)
			{
				Atom const a = atoms[index];
				if (a != 0 && a != InlineHashtable::DELETED)
				{
					return (index>>1)+1;
				}
				index += 2;
			}
		}
		return 0;
	}
}
<commit_msg>comment<commit_after>/* ***** BEGIN LICENSE BLOCK *****
 * Version: MPL 1.1/GPL 2.0/LGPL 2.1
 *
 * The contents of this file are subject to the Mozilla Public License Version
 * 1.1 (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.mozilla.org/MPL/
 *
 * Software distributed under the License is distributed on an "AS IS" basis,
 * WITHOUT WARRANTY OF ANY KIND, either express or implied. See the License
 * for the specific language governing rights and limitations under the
 * License.
 *
 * The Original Code is [Open Source Virtual Machine.].
 *
 * The Initial Developer of the Original Code is
 * Adobe System Incorporated.
 * Portions created by the Initial Developer are Copyright (C) 2004-2006
 * the Initial Developer. All Rights Reserved.
 *
 * Contributor(s):
 *   Adobe AS3 Team
 *
 * Alternatively, the contents of this file may be used under the terms of
 * either the GNU General Public License Version 2 or later (the "GPL"), or
 * the GNU Lesser General Public License Version 2.1 or later (the "LGPL"),
 * in which case the provisions of the GPL or the LGPL are applicable instead
 * of those above. If you wish to allow use of your version of this file only
 * under the terms of either the GPL or the LGPL, and not to allow others to
 * use your version of this file under the terms of the MPL, indicate your
 * decision by deleting the provisions above and replace them with the notice
 * and other provisions required by the GPL or the LGPL. If you do not delete
 * the provisions above, a recipient may use your version of this file under
 * the terms of any one of the MPL, the GPL or the LGPL.
 *
 * ***** END LICENSE BLOCK ***** */

#ifdef AVMSHELL_BUILD
#include "avmshell.h"
#else
// player
#include "platformbuild.h"
#include "avmplayer.h"
#include "DictionaryGlue.h"
#endif

using namespace MMgc;

namespace avmplus
{
	DictionaryClass::DictionaryClass(VTable *vtable)
	: ClassClosure(vtable)
	{
		createVanillaPrototype();
		vtable->traits->itraits->isDictionary = true;
	}

	ScriptObject *DictionaryClass::createInstance(VTable *ivtable, ScriptObject* /*delegate*/)
	{
 		GCAssert(ivtable->traits->isDictionary == true);
		return new (core()->GetGC(), ivtable->getExtraSize()) DictionaryObject(ivtable, prototype);
	}

	DictionaryObject::DictionaryObject(VTable *vtable, ScriptObject *delegate)
		: ScriptObject(vtable, delegate)
	{
		GCAssert(vtable->traits->isDictionary == true);
	}

	void DictionaryObject::init(bool weakKeys)
	{
		GCAssert(vtable->traits->isDictionary == true);
		MMgc::GC* gc = this->gc();

		HeapHashtable* ht = weakKeys ? new (gc) WeakKeyHashtable(gc)
									: new (gc) HeapHashtable(gc);

		//store pointer of newly created hashtable, encapsulated with writebarrier,
		//at the hashtable offset address of the corresponding traits
		union {
			uint8_t* p;
			HeapHashtable** hht;
		};
		p = (uint8_t*)this + vtable->traits->getHashtableOffset();
		WB(gc, this, hht, ht);
	}

	Atom FASTCALL DictionaryObject::getKeyFromObject(Atom key) const
	{
		AvmAssert(AvmCore::isObject(key));
		ScriptObject* obj = AvmCore::atomToScriptObject(key);
		AvmAssert(Traits::getBuiltinType(obj->traits()) != BUILTIN_qName);
		AvmAssert(MMgc::GC::Size(obj) >= sizeof(ScriptObject));
		(void)obj;

		// This commented-out code probably pertains to Bugzilla 507699:
		// "Dictionary key of Xml type are not matching."
		//
		// FIXME: this doesn't work, need to convert back to an XMLObject
		// on the way out or intern XMLObject's somehow
		//if(AvmCore::isXML(key))
		//	key = AvmCore::genericObjectToAtom(AvmCore::atomToXML(key));

		return key;
	}

	Atom DictionaryObject::getAtomProperty(Atom key) const
	{
		if (AvmCore::isObject(key))
		{
			return getHeapHashtable()->get(getKeyFromObject(key));
		}

		return ScriptObject::getAtomProperty(key);
	}

	bool DictionaryObject::hasAtomProperty(Atom key) const
	{
		if (AvmCore::isObject(key))
		{
			return getHeapHashtable()->contains(getKeyFromObject(key));
		}

		return ScriptObject::hasAtomProperty(key);
	}

	bool DictionaryObject::deleteAtomProperty(Atom key)
	{
		if (AvmCore::isObject(key))
		{
			getHeapHashtable()->remove(getKeyFromObject(key));
			return true;
		}

		return ScriptObject::deleteAtomProperty(key);
	}

	void DictionaryObject::setAtomProperty(Atom key, Atom value)
	{
		if (AvmCore::isObject(key))
		{
			getHeapHashtable()->add(getKeyFromObject(key), value);
			return;
		}

		ScriptObject::setAtomProperty(key, value);
	}

	Atom DictionaryObject::nextName(int index)
	{
		Atom k = ScriptObject::nextName(index);

		if (AvmCore::isGenericObject(k) && getHeapHashtable()->weakKeys())
		{
			GCWeakRef* ref = (GCWeakRef*)AvmCore::atomToGenericObject(k);
			union {
				GCObject* key_o;
				ScriptObject* key;
			};
			key_o = ref->get();
			if (key)
			{
				AvmAssert(key->traits() != NULL);
				return key->atom();
			}

			return undefinedAtom;
		}

		return k;
	}

	int DictionaryObject::nextNameIndex(int index)
	{
		AvmAssert(index >= 0);

		// this can happen if you break in debugger in a subclasses constructor before super
		// has been called -- let's do it in all builds, it's better than crashing.
		HeapHashtable* hht = getHeapHashtable();
		if (!hht)
		{
			return 0;
		}

		// Advance to first non-empty slot.
		index <<= 1;
		InlineHashtable* const ht = hht->get_ht();
		Atom* const atoms = ht->getAtoms();
		int const numAtoms = ht->getCapacity();
		if (hht->weakKeys())
		{
			while (index < numAtoms)
			{
				Atom const a = atoms[index];
				if (AvmCore::isGenericObject(a))
				{
					GCWeakRef *weakRef = (GCWeakRef*)AvmCore::atomToGenericObject(a);
					if (weakRef->get())
						return (index>>1)+1;

					atoms[index] = InlineHashtable::DELETED;
					atoms[index+1] = InlineHashtable::DELETED;
					ht->setHasDeletedItems();
				}
				else if (a != 0 && a != InlineHashtable::DELETED)
				{
					return (index>>1)+1;
				}
				index += 2;
			}
		}
		else
		{
			while (index < numAtoms)
			{
				Atom const a = atoms[index];
				if (a != 0 && a != InlineHashtable::DELETED)
				{
					return (index>>1)+1;
				}
				index += 2;
			}
		}
		return 0;
	}
}
<|endoftext|>


<commit_before>#ifndef _MZ_SHARE_RBTREE_HPP_
#define _MZ_SHARE_RBTREE_HPP_

#include "Metazion/Share/ShareInclude.hpp"

DECL_NAMESPACE_MZ_SHARE_BEGIN

template<typename ValueType>
struct RBTreeNode {
    using Value_t = ValueType;

    enum Color{
        RED = 0,
        BLACK,
    };

    RBTreeNode* m_parent = nullptr;
    RBTreeNode* m_right = nullptr;
    RBTreeNode* m_left = nullptr;
    Color m_color = RED;

    Value_t m_value;

    bool IsRed() const {
        return m_color == RED;
    }

    void SetRed() {
        m_color = RED;
    }

    bool IsBlack() const {
        return m_color == BLACK;
    }

    void SetBlack() {
        m_color = BLACK;
    }

    RBTreeNode* Backward() {
        auto node = this;
        if (node->m_parent == node) {
            return nullptr;
        }

        if (!IsNull(node->m_left)) {
            node = node->m_left;
            while (!IsNull(node->m_right)) {
                node = node->m_right;
            }

            return node;
        }

        auto parent = node->m_parent;
        while (!IsNull(parent) && node == parent->m_left) {
            node = parent;
            parent = node->m_parent;
        }

        return parent;
    }

    RBTreeNode* Forward() {
        auto node = this;
        if (node->m_parent == node) {
            return nullptr;
        }

        if (!IsNull(node->m_right)) {
            node = node->m_right;
            while (!IsNull(node->m_left)) {
                node = node->m_left;
            }

            return node;
        }

        auto parent = node->m_parent;
        while (!IsNull(parent) && node == parent->m_right) {
            node = parent;
            parent = node->m_parent;
        }

        return parent;
    }
};

template<typename NodeType
, typename CompareType = LessCompare<typename NodeType::Value_t>
>
class RBTree {
    DISALLOW_COPY_AND_ASSIGN(RBTree)

    using Node_t = NodeType;
    using Compare_t = CompareType;
    using Value_t = typename Node_t::Value_t;

public:
    RBTree()
        : m_root(nullptr)
        , m_size(0) {}

    ~RBTree() {}

public:
    void Clear() {
        m_root = nullptr;
        m_size = 0;
    }

    bool IsEmpty() const {
        return nullptr == m_root;
    }

    int GetSize() const {
        return m_size;
    }

    Node_t* First() {
        auto next = m_root;
        if (IsNull(next)) {
            return nullptr;
        }

        while (!IsNull(next->m_left)) {
            next = next->m_left;
        }

        return next;
    }

    Node_t* Last() {
        auto next = m_root;
        if (IsNull(next)) {
            return nullptr;
        }

        while (!IsNull(next->m_right)) {
            next = next->m_right;
        }

        return next;
    }

    Node_t* InsertUnique(Node_t* node) {
        ASSERT_TRUE(!IsNull(node));

        Node_t* parent = nullptr;
        auto temp = &m_root;

        while (!IsNull(*temp)) {
            parent = *temp;

            const auto ret = m_compare(node->m_value, parent->m_value);
            if (ret < 0) {
                temp = &(*temp)->m_left;
            }
            else if (ret > 0) {
                temp = &(*temp)->m_right;
            }
            else  {
                return nullptr;
            }
        }

        Link(node, parent, *temp);
        InsertColor(node);
        ++m_size;
        return node;
    }

    Node_t* InsertRepeat(Node_t* node) {
        ASSERT_TRUE(!IsNull(node));

        Node_t* parent = nullptr;
        auto temp = &m_root;

        while (!IsNull(*temp)) {
            parent = *temp;

            const auto ret = m_compare(node->m_value, parent->m_value);
            if (ret < 0) {
                temp = &(*temp)->m_left;
            }
            else {
                temp = &(*temp)->m_right;
            }
        }

        Link(node, parent, *temp);
        InsertColor(node);
        ++m_size;
        return node;
    }

    Node_t* Remove(Node_t* node) {
        ASSERT_TRUE(!IsNull(node));

        auto forward = node->Forward();

        Node_t* child = nullptr;
        Node_t* parent = nullptr;
        auto color = Node_t::RED;

        do {
            if (IsNull(node->m_left)) {
                child = node->m_right;
            }
            else if (IsNull(node->m_right)) {
                child = node->m_left;
            }
            else {
                auto old = node;
                node = node->m_right;

                auto left = node->m_left;
                while (!IsNull(left)) {
                    node = left;
                    left = node->m_left;
                }

                auto oldParent = old->m_parent;
                if (!IsNull(oldParent)) {
                    if (oldParent->m_left == old) {
                        oldParent->m_left = node;
                    }
                    else {
                        oldParent->m_right = node;
                    }
                }
                else {
                    m_root = node;
                }

                child = node->m_right;
                parent = node->m_parent;
                color = node->m_color;

                if (parent == old) {
                    parent = node;
                }
                else {
                    if (!IsNull(child)) {
                        child->m_parent = parent;
                    }

                    parent->m_left = child;
                    node->m_right = old->m_right;
                    old->m_right->m_parent = node;
                }

                node->m_color = old->m_color;
                node->m_parent = old->m_parent;
                node->m_left = old->m_left;
                old->m_left->m_parent = node;

                break;
            }

            parent = node->m_parent;
            color = node->m_color;

            if (!IsNull(child)) {
                child->m_parent = parent;
            }

            if (!IsNull(parent)) {
                if (parent->m_left == node) {
                    parent->m_left = child;
                }
                else {
                    parent->m_right = child;
                }
            }
            else {
                m_root = child;
            }
        } while (false);

        if (color == Node_t::BLACK) {
            RemoveColor(child, parent);
        }

        --m_size;
        return forward;
    }

    Node_t* Search(const Value_t& value) {
        auto node = m_root;
        while (!IsNull(node)) {
            const auto ret = m_compare(value, node->m_value);
            if (ret < 0) {
                node = node->m_left;
            }
            else if (ret > 0) {
                node = node->m_right;
            }
            else {
                return node;
            }
        }

        return nullptr;
    }

    Node_t* Replace(Node_t* victim, Node_t* replacement) {
        ASSERT_TRUE(!IsNull(victim));
        ASSERT_TRUE(!IsNull(replacement));

        auto parent = victim->m_parent;
        if (!IsNull(parent)) {
            if (victim == parent->m_left) {
                parent->m_left = replacement;
            }
            else {
                parent->m_right = replacement;
            }
        }
        else {
            m_root = replacement;
        }

        if (!IsNull(victim->m_left)) {
            victim->m_left->m_parent = replacement;
        }

        if (!IsNull(victim->m_right)) {
            victim->m_right->m_parent = replacement;
        }

        *replacement = *victim;
        return victim;
    }

    Node_t* Backward(Node_t* node) {
        ASSERT_TRUE(!IsNull(node));

        return node->Backward();
    }

    Node_t* Forward(Node_t* node) {
        ASSERT_TRUE(!IsNull(node));

        return node->Forward();
    }

private:

    void Link(Node_t* node, Node_t* parent, Node_t*& link) {
        node->m_parent = parent;
        node->m_left = node->m_right = nullptr;

        link = node;
    }

    void InsertColor(Node_t* node) {
        auto parent = node->m_parent;
        while (!IsNull(parent) && parent->IsRed()) {
            auto gparent = parent->m_parent;
            if (parent == gparent->m_left) {
                auto uncle = gparent->m_right;
                if (uncle && uncle->IsRed()) {
                    uncle->SetBlack();
                    parent->SetBlack();
                    gparent->SetRed();
                    node = gparent;
                    continue;
                }

                if (parent->m_right == node) {
                    RotateLeft(parent);
                    auto temp = parent;
                    parent = node;
                    node = temp;
                }

                parent->SetBlack();
                gparent->SetRed();
                RotateRight(gparent);
            }
            else {
                auto uncle = gparent->m_left;
                if (uncle && uncle->IsRed()) {
                    uncle->SetBlack();
                    parent->SetBlack();
                    gparent->SetRed();
                    node = gparent;
                    continue;
                }

                if (parent->m_left == node) {
                    RotateRight(parent);
                    auto temp = parent;
                    parent = node;
                    node = temp;
                }

                parent->SetBlack();
                gparent->SetRed();
                RotateLeft(gparent);
            }

            parent = node->m_parent;
        }

        m_root->SetBlack();
    }

    void RemoveColor(Node_t* node, Node_t* parent) {
        while ((IsNull(node) || node->IsBlack()) && node != m_root) {
            if (parent->m_left == node) {
                auto other = parent->m_right;
                if (other->IsRed()) {
                    other->SetBlack();
                    parent->SetRed();
                    RotateLeft(parent);
                    other = parent->m_right;
                }

                if ((IsNull(other->m_left) || other->m_left->IsBlack()) &&
                    (IsNull(other->m_right) || other->m_right->IsBlack())) {
                    other->SetRed();
                    node = parent;
                    parent = node->m_parent;
                }
                else {
                    if (IsNull(other->m_right) || other->m_right->IsBlack()) {
                        other->m_left->SetBlack();
                        other->SetRed();
                        RotateRight(other);
                        other = parent->m_right;
                    }

                    other->m_color = parent->m_color;
                    parent->SetBlack();
                    other->m_right->SetBlack();
                    RotateLeft(parent);
                    node = m_root;
                    break;
                }
            }
            else {
                auto other = parent->m_left;
                if (other->IsRed()) {
                    other->SetBlack();
                    parent->SetRed();
                    RotateRight(parent);
                    other = parent->m_left;
                }

                if ((IsNull(other->m_left) || other->m_left->IsBlack()) &&
                    (IsNull(other->m_right) || other->m_right->IsBlack())) {
                    other->SetRed();
                    node = parent;
                    parent = node->m_parent;
                }
                else {
                    if (IsNull(other->m_left) || other->m_left->IsBlack()) {
                        other->m_right->SetBlack();
                        other->SetRed();
                        RotateLeft(other);
                        other = parent->m_left;
                    }

                    other->m_color = parent->m_color;
                    parent->SetBlack();
                    other->m_left->SetBlack();
                    RotateRight(parent);
                    node = m_root;
                    break;
                }
            }
        }

        if (!IsNull(node)) {
            node->SetBlack();
        }
    }

    void RotateLeft(Node_t* node) {
        auto right = node->m_right;
        auto parent = node->m_parent;
        node->m_right = right->m_left;

        if (!IsNull(right->m_left)) {
            right->m_left->m_parent = node;
        }

        right->m_left = node;
        right->m_parent = parent;

        if (!IsNull(parent)) {
            if (node == parent->m_left) {
                parent->m_left = right;
            }
            else {
                parent->m_right = right;
            }
        }
        else {
            m_root = right;
        }

        node->m_parent = right;
    }

    void RotateRight(Node_t* node) {
        auto left = node->m_left;
        auto parent = node->m_parent;
        node->m_left = left->m_right;

        if (!IsNull(left->m_right)) {
            left->m_right->m_parent = node;
        }

        left->m_right = node;
        left->m_parent = parent;

        if (!IsNull(parent)) {
            if (node == parent->m_right) {
                parent->m_right = left;
            }
            else {
                parent->m_left = left;
            }
        }
        else {
            m_root = left;
        }

        node->m_parent = left;
    }

private:
    Compare_t m_compare;
    Node_t* m_root;
    int m_size;
};

DECL_NAMESPACE_MZ_SHARE_END

#endif // _MZ_SHARE_RBTREE_HPP_
<commit_msg>Some trivial modifications.<commit_after>#ifndef _MZ_SHARE_RBTREE_HPP_
#define _MZ_SHARE_RBTREE_HPP_

#include "Metazion/Share/ShareInclude.hpp"

DECL_NAMESPACE_MZ_SHARE_BEGIN

template<typename ValueType>
struct RBTreeNode {
    using Value_t = ValueType;

    enum Color {
        RED = 0,
        BLACK,
    };

    RBTreeNode* m_parent = nullptr;
    RBTreeNode* m_right = nullptr;
    RBTreeNode* m_left = nullptr;
    Color m_color = RED;

    Value_t m_value;

    bool IsRed() const {
        return m_color == RED;
    }

    void SetRed() {
        m_color = RED;
    }

    bool IsBlack() const {
        return m_color == BLACK;
    }

    void SetBlack() {
        m_color = BLACK;
    }

    RBTreeNode* Backward() {
        auto node = this;
        if (node->m_parent == node) {
            return nullptr;
        }

        if (!IsNull(node->m_left)) {
            node = node->m_left;
            while (!IsNull(node->m_right)) {
                node = node->m_right;
            }

            return node;
        }

        auto parent = node->m_parent;
        while (!IsNull(parent) && node == parent->m_left) {
            node = parent;
            parent = node->m_parent;
        }

        return parent;
    }

    RBTreeNode* Forward() {
        auto node = this;
        if (node->m_parent == node) {
            return nullptr;
        }

        if (!IsNull(node->m_right)) {
            node = node->m_right;
            while (!IsNull(node->m_left)) {
                node = node->m_left;
            }

            return node;
        }

        auto parent = node->m_parent;
        while (!IsNull(parent) && node == parent->m_right) {
            node = parent;
            parent = node->m_parent;
        }

        return parent;
    }
};

template<typename NodeType
, typename CompareType = LessCompare<typename NodeType::Value_t>
>
class RBTree {
    DISALLOW_COPY_AND_ASSIGN(RBTree)

    using Node_t = NodeType;
    using Compare_t = CompareType;
    using Value_t = typename Node_t::Value_t;

public:
    RBTree()
        : m_root(nullptr)
        , m_size(0) {}

    ~RBTree() {}

public:
    void Clear() {
        m_root = nullptr;
        m_size = 0;
    }

    bool IsEmpty() const {
        return nullptr == m_root;
    }

    int GetSize() const {
        return m_size;
    }

    Node_t* First() {
        auto next = m_root;
        if (IsNull(next)) {
            return nullptr;
        }

        while (!IsNull(next->m_left)) {
            next = next->m_left;
        }

        return next;
    }

    Node_t* Last() {
        auto next = m_root;
        if (IsNull(next)) {
            return nullptr;
        }

        while (!IsNull(next->m_right)) {
            next = next->m_right;
        }

        return next;
    }

    Node_t* InsertUnique(Node_t* node) {
        ASSERT_TRUE(!IsNull(node));

        Node_t* parent = nullptr;
        auto temp = &m_root;

        while (!IsNull(*temp)) {
            parent = *temp;

            const auto ret = m_compare(node->m_value, parent->m_value);
            if (ret < 0) {
                temp = &(*temp)->m_left;
            }
            else if (ret > 0) {
                temp = &(*temp)->m_right;
            }
            else  {
                return nullptr;
            }
        }

        Link(node, parent, *temp);
        InsertColor(node);
        ++m_size;
        return node;
    }

    Node_t* InsertRepeat(Node_t* node) {
        ASSERT_TRUE(!IsNull(node));

        Node_t* parent = nullptr;
        auto temp = &m_root;

        while (!IsNull(*temp)) {
            parent = *temp;

            const auto ret = m_compare(node->m_value, parent->m_value);
            if (ret < 0) {
                temp = &(*temp)->m_left;
            }
            else {
                temp = &(*temp)->m_right;
            }
        }

        Link(node, parent, *temp);
        InsertColor(node);
        ++m_size;
        return node;
    }

    Node_t* Remove(Node_t* node) {
        ASSERT_TRUE(!IsNull(node));

        auto forward = node->Forward();

        Node_t* child = nullptr;
        Node_t* parent = nullptr;
        auto color = Node_t::RED;

        do {
            if (IsNull(node->m_left)) {
                child = node->m_right;
            }
            else if (IsNull(node->m_right)) {
                child = node->m_left;
            }
            else {
                auto old = node;
                node = node->m_right;

                auto left = node->m_left;
                while (!IsNull(left)) {
                    node = left;
                    left = node->m_left;
                }

                auto oldParent = old->m_parent;
                if (!IsNull(oldParent)) {
                    if (oldParent->m_left == old) {
                        oldParent->m_left = node;
                    }
                    else {
                        oldParent->m_right = node;
                    }
                }
                else {
                    m_root = node;
                }

                child = node->m_right;
                parent = node->m_parent;
                color = node->m_color;

                if (parent == old) {
                    parent = node;
                }
                else {
                    if (!IsNull(child)) {
                        child->m_parent = parent;
                    }

                    parent->m_left = child;
                    node->m_right = old->m_right;
                    old->m_right->m_parent = node;
                }

                node->m_color = old->m_color;
                node->m_parent = old->m_parent;
                node->m_left = old->m_left;
                old->m_left->m_parent = node;

                break;
            }

            parent = node->m_parent;
            color = node->m_color;

            if (!IsNull(child)) {
                child->m_parent = parent;
            }

            if (!IsNull(parent)) {
                if (parent->m_left == node) {
                    parent->m_left = child;
                }
                else {
                    parent->m_right = child;
                }
            }
            else {
                m_root = child;
            }
        } while (false);

        if (color == Node_t::BLACK) {
            RemoveColor(child, parent);
        }

        --m_size;
        return forward;
    }

    Node_t* Search(const Value_t& value) {
        auto node = m_root;
        while (!IsNull(node)) {
            const auto ret = m_compare(value, node->m_value);
            if (ret < 0) {
                node = node->m_left;
            }
            else if (ret > 0) {
                node = node->m_right;
            }
            else {
                return node;
            }
        }

        return nullptr;
    }

    Node_t* Replace(Node_t* victim, Node_t* replacement) {
        ASSERT_TRUE(!IsNull(victim));
        ASSERT_TRUE(!IsNull(replacement));

        auto parent = victim->m_parent;
        if (!IsNull(parent)) {
            if (victim == parent->m_left) {
                parent->m_left = replacement;
            }
            else {
                parent->m_right = replacement;
            }
        }
        else {
            m_root = replacement;
        }

        if (!IsNull(victim->m_left)) {
            victim->m_left->m_parent = replacement;
        }

        if (!IsNull(victim->m_right)) {
            victim->m_right->m_parent = replacement;
        }

        *replacement = *victim;
        return victim;
    }

    Node_t* Backward(Node_t* node) {
        ASSERT_TRUE(!IsNull(node));

        return node->Backward();
    }

    Node_t* Forward(Node_t* node) {
        ASSERT_TRUE(!IsNull(node));

        return node->Forward();
    }

private:

    void Link(Node_t* node, Node_t* parent, Node_t*& link) {
        node->m_parent = parent;
        node->m_left = node->m_right = nullptr;

        link = node;
    }

    void InsertColor(Node_t* node) {
        auto parent = node->m_parent;
        while (!IsNull(parent) && parent->IsRed()) {
            auto gparent = parent->m_parent;
            if (parent == gparent->m_left) {
                auto uncle = gparent->m_right;
                if (uncle && uncle->IsRed()) {
                    uncle->SetBlack();
                    parent->SetBlack();
                    gparent->SetRed();
                    node = gparent;
                    continue;
                }

                if (parent->m_right == node) {
                    RotateLeft(parent);
                    auto temp = parent;
                    parent = node;
                    node = temp;
                }

                parent->SetBlack();
                gparent->SetRed();
                RotateRight(gparent);
            }
            else {
                auto uncle = gparent->m_left;
                if (uncle && uncle->IsRed()) {
                    uncle->SetBlack();
                    parent->SetBlack();
                    gparent->SetRed();
                    node = gparent;
                    continue;
                }

                if (parent->m_left == node) {
                    RotateRight(parent);
                    auto temp = parent;
                    parent = node;
                    node = temp;
                }

                parent->SetBlack();
                gparent->SetRed();
                RotateLeft(gparent);
            }

            parent = node->m_parent;
        }

        m_root->SetBlack();
    }

    void RemoveColor(Node_t* node, Node_t* parent) {
        while ((IsNull(node) || node->IsBlack()) && node != m_root) {
            if (parent->m_left == node) {
                auto other = parent->m_right;
                if (other->IsRed()) {
                    other->SetBlack();
                    parent->SetRed();
                    RotateLeft(parent);
                    other = parent->m_right;
                }

                if ((IsNull(other->m_left) || other->m_left->IsBlack()) &&
                    (IsNull(other->m_right) || other->m_right->IsBlack())) {
                    other->SetRed();
                    node = parent;
                    parent = node->m_parent;
                }
                else {
                    if (IsNull(other->m_right) || other->m_right->IsBlack()) {
                        other->m_left->SetBlack();
                        other->SetRed();
                        RotateRight(other);
                        other = parent->m_right;
                    }

                    other->m_color = parent->m_color;
                    parent->SetBlack();
                    other->m_right->SetBlack();
                    RotateLeft(parent);
                    node = m_root;
                    break;
                }
            }
            else {
                auto other = parent->m_left;
                if (other->IsRed()) {
                    other->SetBlack();
                    parent->SetRed();
                    RotateRight(parent);
                    other = parent->m_left;
                }

                if ((IsNull(other->m_left) || other->m_left->IsBlack()) &&
                    (IsNull(other->m_right) || other->m_right->IsBlack())) {
                    other->SetRed();
                    node = parent;
                    parent = node->m_parent;
                }
                else {
                    if (IsNull(other->m_left) || other->m_left->IsBlack()) {
                        other->m_right->SetBlack();
                        other->SetRed();
                        RotateLeft(other);
                        other = parent->m_left;
                    }

                    other->m_color = parent->m_color;
                    parent->SetBlack();
                    other->m_left->SetBlack();
                    RotateRight(parent);
                    node = m_root;
                    break;
                }
            }
        }

        if (!IsNull(node)) {
            node->SetBlack();
        }
    }

    void RotateLeft(Node_t* node) {
        auto right = node->m_right;
        auto parent = node->m_parent;
        node->m_right = right->m_left;

        if (!IsNull(right->m_left)) {
            right->m_left->m_parent = node;
        }

        right->m_left = node;
        right->m_parent = parent;

        if (!IsNull(parent)) {
            if (node == parent->m_left) {
                parent->m_left = right;
            }
            else {
                parent->m_right = right;
            }
        }
        else {
            m_root = right;
        }

        node->m_parent = right;
    }

    void RotateRight(Node_t* node) {
        auto left = node->m_left;
        auto parent = node->m_parent;
        node->m_left = left->m_right;

        if (!IsNull(left->m_right)) {
            left->m_right->m_parent = node;
        }

        left->m_right = node;
        left->m_parent = parent;

        if (!IsNull(parent)) {
            if (node == parent->m_right) {
                parent->m_right = left;
            }
            else {
                parent->m_left = left;
            }
        }
        else {
            m_root = left;
        }

        node->m_parent = left;
    }

private:
    Compare_t m_compare;
    Node_t* m_root;
    int m_size;
};

DECL_NAMESPACE_MZ_SHARE_END

#endif // _MZ_SHARE_RBTREE_HPP_
<|endoftext|>


<commit_before>#include "editor/editor_notes.hpp"

#include "platform/platform.hpp"

#include "coding/internal/file_data.hpp"

#include "geometry/mercator.hpp"

#include "base/string_utils.hpp"
#include "base/assert.hpp"
#include "base/logging.hpp"
#include "base/timer.hpp"

#include "std/future.hpp"

#include "3party/pugixml/src/pugixml.hpp"

namespace
{
bool LoadFromXml(pugi::xml_document const & xml,
                 vector<editor::Note> & notes,
                 uint32_t & uploadedNotesCount)
{
  uint64_t notesCount;
  auto const root = xml.child("notes");
  if (!strings::to_uint64(root.attribute("count").value(), notesCount))
    return false;
  uploadedNotesCount = static_cast<uint32_t>(notesCount);
  for (auto const xNode : root.select_nodes("note"))
  {
    m2::PointD point;

    auto const node = xNode.node();
    auto const point_x = node.attribute("x");
    if (!point_x || !strings::to_double(point_x.value(), point.x))
      return false;

    auto const point_y = node.attribute("y");
    if (!point_y || !strings::to_double(point_y.value(), point.y))
      return false;

    auto const text = node.attribute("text");
    if (!text)
      return false;

    notes.emplace_back(point, text.value());
  }
  return true;
}

void SaveToXml(vector<editor::Note> const & notes,
               pugi::xml_document & xml,
               uint32_t const UploadedNotesCount)
{
  auto root = xml.append_child("notes");
  root.append_attribute("count") = UploadedNotesCount;
  for (auto const & note : notes)
  {
    auto node = root.append_child("note");
    node.append_attribute("x") = DebugPrint(note.m_point.x).data();
    node.append_attribute("y") = DebugPrint(note.m_point.y).data();
    node.append_attribute("text") = note.m_note.data();
  }
}

vector<editor::Note> MoveNoteVectorAtomicly(vector<editor::Note> && notes, mutex & mu)
{
  lock_guard<mutex> g(mu);
  return move(notes);
}
}  // namespace

namespace editor
{
shared_ptr<Notes> Notes::MakeNotes(string const & fileName)
{
  return shared_ptr<Notes>(new Notes(fileName));
}

Notes::Notes(string const & fileName)
    : m_fileName(fileName)
{
  Load();
}

void Notes::CreateNote(m2::PointD const & point, string const & text)
{
  lock_guard<mutex> g(m_mu);
  m_notes.emplace_back(point, text);
  Save();
}

void Notes::Upload(osm::OsmOAuth const & auth)
{
  auto const launch = [this, &auth]()
  {
    // Capture self to keep it from destruction until this thread is done.
    auto const self = shared_from_this();
    return async(launch::async,
                 [self, auth]()
                 {
                   auto const notes = MoveNoteVectorAtomicly(move(self->m_notes),
                                                             self->m_mu);

                   vector<Note> unuploaded;
                   osm::ServerApi06 api(auth);
                   for (auto const & note : notes)
                   {
                     try
                     {
                       api.CreateNote(MercatorBounds::ToLatLon(note.m_point), note.m_note);
                       ++self->m_uploadedNotes;
                     }
                     catch (osm::ServerApi06::ServerApi06Exception const & e)
                     {
                       LOG(LERROR, ("Can't upload note.", e.Msg()));
                       unuploaded.push_back(note);
                     }
                   }

                   lock_guard<mutex> g(self->m_mu);
                   self->m_notes.insert(end(self->m_notes), begin(unuploaded), end(unuploaded));
                 });
  };

  // Do not run more than one upload thread at a time.
  static auto future = launch();
  auto const status = future.wait_for(milliseconds(0));
  if (status == future_status::ready)
    future = launch();
}

bool Notes::Load()
{
  string content;
  try
  {
    auto const reader = GetPlatform().GetReader(m_fileName);
    reader->ReadAsString(content);
  }
  catch (FileAbsentException const &)
  {
    LOG(LINFO, ("No edits file."));
    return true;
  }
  catch (Reader::Exception const &)
  {
    LOG(LERROR, ("Can't process file.", m_fileName));
    return false;
  }

  pugi::xml_document xml;
  if (!xml.load_buffer(content.data(), content.size()))
  {
    LOG(LERROR, ("Can't load notes, xml is illformed."));
    return false;
  }

  lock_guard<mutex> g(m_mu);
  m_notes.clear();
  if (!LoadFromXml(xml, m_notes, m_uploadedNotes))
  {
    LOG(LERROR, ("Can't load notes, file is illformed."));
    return false;
  }

  return true;
}

/// Not thread-safe, use syncronization.
bool Notes::Save()
{
  pugi::xml_document xml;
  SaveToXml(m_notes, xml, m_uploadedNotes);

  string const tmpFileName = m_fileName + ".tmp";
  if (!xml.save_file(tmpFileName.data(), "  "))
  {
    LOG(LERROR, ("Can't save map edits into", tmpFileName));
    return false;
  }
  else if (!my::RenameFileX(tmpFileName, m_fileName))
  {
    LOG(LERROR, ("Can't rename file", tmpFileName, "to", m_fileName));
    return false;
  }
  return true;
}
}
<commit_msg>[editor] Call Notes::Save from Notes::Upload.<commit_after>#include "editor/editor_notes.hpp"

#include "platform/platform.hpp"

#include "coding/internal/file_data.hpp"

#include "geometry/mercator.hpp"

#include "base/string_utils.hpp"
#include "base/assert.hpp"
#include "base/logging.hpp"
#include "base/timer.hpp"

#include "std/future.hpp"

#include "3party/pugixml/src/pugixml.hpp"

namespace
{
bool LoadFromXml(pugi::xml_document const & xml,
                 vector<editor::Note> & notes,
                 uint32_t & uploadedNotesCount)
{
  uint64_t notesCount;
  auto const root = xml.child("notes");
  if (!strings::to_uint64(root.attribute("count").value(), notesCount))
    return false;
  uploadedNotesCount = static_cast<uint32_t>(notesCount);
  for (auto const xNode : root.select_nodes("note"))
  {
    m2::PointD point;

    auto const node = xNode.node();
    auto const point_x = node.attribute("x");
    if (!point_x || !strings::to_double(point_x.value(), point.x))
      return false;

    auto const point_y = node.attribute("y");
    if (!point_y || !strings::to_double(point_y.value(), point.y))
      return false;

    auto const text = node.attribute("text");
    if (!text)
      return false;

    notes.emplace_back(point, text.value());
  }
  return true;
}

void SaveToXml(vector<editor::Note> const & notes,
               pugi::xml_document & xml,
               uint32_t const UploadedNotesCount)
{
  auto root = xml.append_child("notes");
  root.append_attribute("count") = UploadedNotesCount;
  for (auto const & note : notes)
  {
    auto node = root.append_child("note");
    node.append_attribute("x") = DebugPrint(note.m_point.x).data();
    node.append_attribute("y") = DebugPrint(note.m_point.y).data();
    node.append_attribute("text") = note.m_note.data();
  }
}

vector<editor::Note> MoveNoteVectorAtomicly(vector<editor::Note> && notes, mutex & mu)
{
  lock_guard<mutex> g(mu);
  return move(notes);
}
}  // namespace

namespace editor
{
shared_ptr<Notes> Notes::MakeNotes(string const & fileName)
{
  return shared_ptr<Notes>(new Notes(fileName));
}

Notes::Notes(string const & fileName)
    : m_fileName(fileName)
{
  Load();
}

void Notes::CreateNote(m2::PointD const & point, string const & text)
{
  lock_guard<mutex> g(m_mu);
  m_notes.emplace_back(point, text);
  Save();
}

void Notes::Upload(osm::OsmOAuth const & auth)
{
  auto const launch = [this, &auth]()
  {
    // Capture self to keep it from destruction until this thread is done.
    auto const self = shared_from_this();
    return async(launch::async,
                 [self, auth]()
                 {
                   auto const notes = MoveNoteVectorAtomicly(move(self->m_notes),
                                                             self->m_mu);

                   vector<Note> unuploaded;
                   osm::ServerApi06 api(auth);
                   for (auto const & note : notes)
                   {
                     try
                     {
                       api.CreateNote(MercatorBounds::ToLatLon(note.m_point), note.m_note);
                       ++self->m_uploadedNotes;
                     }
                     catch (osm::ServerApi06::ServerApi06Exception const & e)
                     {
                       LOG(LERROR, ("Can't upload note.", e.Msg()));
                       unuploaded.push_back(note);
                     }
                   }

                   lock_guard<mutex> g(self->m_mu);
                   self->m_notes.insert(end(self->m_notes), begin(unuploaded), end(unuploaded));
                   self->Save();
                 });
  };

  // Do not run more than one upload thread at a time.
  static auto future = launch();
  auto const status = future.wait_for(milliseconds(0));
  if (status == future_status::ready)
    future = launch();
}

bool Notes::Load()
{
  string content;
  try
  {
    auto const reader = GetPlatform().GetReader(m_fileName);
    reader->ReadAsString(content);
  }
  catch (FileAbsentException const &)
  {
    LOG(LINFO, ("No edits file."));
    return true;
  }
  catch (Reader::Exception const &)
  {
    LOG(LERROR, ("Can't process file.", m_fileName));
    return false;
  }

  pugi::xml_document xml;
  if (!xml.load_buffer(content.data(), content.size()))
  {
    LOG(LERROR, ("Can't load notes, xml is illformed."));
    return false;
  }

  lock_guard<mutex> g(m_mu);
  m_notes.clear();
  if (!LoadFromXml(xml, m_notes, m_uploadedNotes))
  {
    LOG(LERROR, ("Can't load notes, file is illformed."));
    return false;
  }

  return true;
}

/// Not thread-safe, use syncronization.
bool Notes::Save()
{
  pugi::xml_document xml;
  SaveToXml(m_notes, xml, m_uploadedNotes);

  string const tmpFileName = m_fileName + ".tmp";
  if (!xml.save_file(tmpFileName.data(), "  "))
  {
    LOG(LERROR, ("Can't save map edits into", tmpFileName));
    return false;
  }
  else if (!my::RenameFileX(tmpFileName, m_fileName))
  {
    LOG(LERROR, ("Can't rename file", tmpFileName, "to", m_fileName));
    return false;
  }
  return true;
}
}
<|endoftext|>


<commit_before>/* Copyright 2018 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/compiler/tf2tensorrt/convert/convert_graph.h"

#include <regex>  // NOLINT

#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include "tensorflow/cc/framework/ops.h"
#include "tensorflow/cc/framework/scope.h"
#include "tensorflow/cc/ops/standard_ops.h"
#include "tensorflow/compiler/tf2tensorrt/convert/convert_nodes.h"
#include "tensorflow/core/common_runtime/device_mgr.h"
#include "tensorflow/core/common_runtime/device_set.h"
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/grappler/clusters/cluster.h"
#include "tensorflow/core/grappler/grappler_item.h"
#include "tensorflow/core/lib/core/status.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/protobuf/config.pb.h"  // NOLINT
#include "tensorflow/core/public/session.h"

#if GOOGLE_CUDA && GOOGLE_TENSORRT

namespace tensorflow {
namespace tensorrt {
namespace convert {

// TODO(laigd): put this into some test utils file.
void ExpectStatus(Status status, error::Code code = error::OK,
                  const char* substr = nullptr) {
  EXPECT_EQ(code, status.code())
      << status << " vs expected error code \"" << error::Code_Name(code)
      << "\" and message \"" << substr << "\"";
  if (substr) {
    EXPECT_THAT(status.error_message(), ::testing::HasSubstr(substr)) << status;
  }
}

class FakeCluster : public grappler::Cluster {
 public:
  FakeCluster() : Cluster(0) {}

  void SetDeviceSet(const DeviceSet* device_set) { device_set_ = device_set; }

  const DeviceSet* GetDeviceSet() const override { return device_set_; }

  string type() const override { return ""; }
  Status Provision() override { return Status::OK(); }
  Status Initialize(const grappler::GrapplerItem& item) override {
    return Status::OK();
  }
  Status Run(const GraphDef& graph_def,
             const std::vector<std::pair<string, Tensor>>& feed,
             const std::vector<string>& fetch, RunMetadata* metadata) override {
    return Status::OK();
  }

 private:
  const DeviceSet* device_set_;
};

TEST(ConvertGraphTest, GetDeviceAndAllocator) {
  ConversionParams params;
  EngineInfo engine_info;
  {
    // params.cluster is not set, and no gpu device is available.
    auto result = GetDeviceAndAllocator(params, engine_info);
    EXPECT_EQ(-1, result.first);
    EXPECT_EQ(nullptr, result.second);
  }

  // Create a session with two (virtual) gpu device.
  SessionOptions options;
  ConfigProto* config = &options.config;
  GPUOptions* gpu_options = config->mutable_gpu_options();
  auto virtual_devices =
      gpu_options->mutable_experimental()->add_virtual_devices();
  virtual_devices->add_memory_limit_mb(200);
  virtual_devices->add_memory_limit_mb(200);
  std::unique_ptr<Session> session(NewSession(options));

  {
    // params.cluster is not set, should find and return first gpu id and
    // corresponding allocator.
    auto result = GetDeviceAndAllocator(params, engine_info);
    EXPECT_EQ(0, result.first);
    EXPECT_NE(nullptr, result.second);
    EXPECT_EQ("GPU_0_bfc", result.second->Name());
  }

  FakeCluster cluster;
  params.cluster = &cluster;
  {
    // params.cluster->GetDeviceSet() returns null, should find and return first
    // gpu id and corresponding allocator.
    auto result = GetDeviceAndAllocator(params, engine_info);
    EXPECT_EQ(0, result.first);
    EXPECT_NE(nullptr, result.second);
    EXPECT_EQ("GPU_0_bfc", result.second->Name());
  }

  // Build the DeviceSet.
  DeviceSet device_set;
  const DeviceMgr* device_mgr = nullptr;
  TF_ASSERT_OK(session->LocalDeviceManager(&device_mgr));
  for (auto d : device_mgr->ListDevices()) {
    device_set.AddDevice(d);
  }
  cluster.SetDeviceSet(&device_set);
  {
    // engine_info.device is not set, should find and return first gpu id and
    // corresponding allocator.
    auto result = GetDeviceAndAllocator(params, engine_info);
    EXPECT_EQ(0, result.first);
    EXPECT_NE(nullptr, result.second);
    EXPECT_EQ("GPU_0_bfc", result.second->Name());
  }

  engine_info.device = "/GPU:1";
  {
    // Set to use second device.
    auto result = GetDeviceAndAllocator(params, engine_info);
    EXPECT_EQ(0, result.first);
    EXPECT_NE(nullptr, result.second);
    EXPECT_EQ("GPU_1_bfc", result.second->Name());
  }

  engine_info.device = "/GPU:3";
  {
    // Set to use nonexistent device.
    auto result = GetDeviceAndAllocator(params, engine_info);
    EXPECT_EQ(-1, result.first);
    EXPECT_EQ(nullptr, result.second);
  }
}

class ConvertAfterShapesTest : public ::testing::Test {
 public:
  Status RunConvertAfterShape(Scope s, GraphDef* output_graph_def) {
    // Create GraphProperties.
    grappler::GrapplerItem item;
    TF_EXPECT_OK(s.ToGraphDef(&item.graph));
    grappler::GraphProperties graph_properties(item);
    TF_EXPECT_OK(graph_properties.InferStatically(true));

    // Construct ConversionParams.
    const std::vector<string> output_names{"output"};
    ConversionParams params;
    params.output_names = &output_names;
    params.max_workspace_size_bytes = 8 << 20;
    params.output_graph_def = output_graph_def;
    params.minimum_segment_size = 1;
    params.grappler_item = &item;
    params.use_calibration = false;
    params.trt_logger_name = "DefaultLogger";

    return ConvertAfterShapes(params);
  }
};

TEST_F(ConvertAfterShapesTest, DirectlyConnectedEngines) {
  // Create the graph. There will be two TRTEngineOps after the conversion, and
  // the upstream TRTEngineOp will have two output connections from the same
  // node:port inside the op to the downstream TRTEngineOp. Then, if it adds the
  // downstream TRTEngineOp first, when adding the upstream op it'll need to
  // update the same output connection twice. This test ensures the correctness
  // of the conversion under such condition.
  Scope s = Scope::NewRootScope();
  auto input = ops::Placeholder(s.WithOpName("input"), DT_FLOAT,
                                ops::Placeholder::Shape({2, 1}));
  // We purposefully choose the name of the root node of each segment, so it'll
  // process the segment in the downstream first, then, when it tries to update
  // the edge between the two TRTEngineOps, it'll try to add the same edge
  // multiple times.
  auto segment_root_1 = ops::Identity(s.WithOpName("segment_root_b"), input);
  auto add1 = ops::Add(s.WithOpName("add1"), segment_root_1, segment_root_1);
  // Add incompatible reshapes that change the batch dimension.
  auto incompatible =
      ops::Reshape(s.WithOpName("reshape1"), add1, Input({1, 2}));
  incompatible =
      ops::Reshape(s.WithOpName("reshape2"), incompatible, Input({2, 1}));

  auto add2 = ops::Add(s.WithOpName("add2"), incompatible, add1);
  auto segment_root_2 = ops::Identity(s.WithOpName("segment_root_a"), add1);
  auto add3 = ops::Add(s.WithOpName("add3"), add2, segment_root_2);
  ops::Identity(s.WithOpName("output"), add3);

  GraphDef output_graph_def;
  TF_EXPECT_OK(RunConvertAfterShape(s, &output_graph_def));

  auto remove_graph_sequence_number = [](std::string node_name) {
    const std::regex pattern("TRTEngineOp_[0-9]+_");
    return std::regex_replace(node_name, pattern, "TRTEngineOp_");
  };
  int num_trt_ops = 0;
  for (const NodeDef& node : output_graph_def.node()) {
    std::string node_name = node.name();
    if (node.op() != "TRTEngineOp") continue;
    node_name = remove_graph_sequence_number(node_name);
    if (node_name == "TRTEngineOp_1") {
      EXPECT_EQ(1, node.input_size());
      EXPECT_EQ("input", node.input(0));
      ++num_trt_ops;
    } else if (node_name == "TRTEngineOp_0") {
      EXPECT_EQ(2, node.input_size());
      EXPECT_EQ("TRTEngineOp_1", remove_graph_sequence_number(node.input(0)));
      EXPECT_EQ("reshape2", node.input(1));
      ++num_trt_ops;
    }
  }
  EXPECT_EQ(2, num_trt_ops);
}

}  // namespace convert
}  // namespace tensorrt
}  // namespace tensorflow

#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
<commit_msg>[TF:TRT] Fix an uninitialized access.<commit_after>/* Copyright 2018 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/compiler/tf2tensorrt/convert/convert_graph.h"

#include <regex>  // NOLINT

#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include "tensorflow/cc/framework/ops.h"
#include "tensorflow/cc/framework/scope.h"
#include "tensorflow/cc/ops/standard_ops.h"
#include "tensorflow/compiler/tf2tensorrt/convert/convert_nodes.h"
#include "tensorflow/core/common_runtime/device_mgr.h"
#include "tensorflow/core/common_runtime/device_set.h"
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/grappler/clusters/cluster.h"
#include "tensorflow/core/grappler/grappler_item.h"
#include "tensorflow/core/lib/core/status.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/protobuf/config.pb.h"  // NOLINT
#include "tensorflow/core/public/session.h"

#if GOOGLE_CUDA && GOOGLE_TENSORRT

namespace tensorflow {
namespace tensorrt {
namespace convert {

// TODO(laigd): put this into some test utils file.
void ExpectStatus(Status status, error::Code code = error::OK,
                  const char* substr = nullptr) {
  EXPECT_EQ(code, status.code())
      << status << " vs expected error code \"" << error::Code_Name(code)
      << "\" and message \"" << substr << "\"";
  if (substr) {
    EXPECT_THAT(status.error_message(), ::testing::HasSubstr(substr)) << status;
  }
}

class FakeCluster : public grappler::Cluster {
 public:
  FakeCluster() : Cluster(0) {}

  void SetDeviceSet(const DeviceSet* device_set) { device_set_ = device_set; }

  const DeviceSet* GetDeviceSet() const override { return device_set_; }

  string type() const override { return ""; }
  Status Provision() override { return Status::OK(); }
  Status Initialize(const grappler::GrapplerItem& item) override {
    return Status::OK();
  }
  Status Run(const GraphDef& graph_def,
             const std::vector<std::pair<string, Tensor>>& feed,
             const std::vector<string>& fetch, RunMetadata* metadata) override {
    return Status::OK();
  }

 private:
  const DeviceSet* device_set_ = nullptr;
};

TEST(ConvertGraphTest, GetDeviceAndAllocator) {
  ConversionParams params;
  EngineInfo engine_info;
  {
    // params.cluster is not set, and no gpu device is available.
    auto result = GetDeviceAndAllocator(params, engine_info);
    EXPECT_EQ(-1, result.first);
    EXPECT_EQ(nullptr, result.second);
  }

  // Create a session with two (virtual) gpu device.
  SessionOptions options;
  ConfigProto* config = &options.config;
  GPUOptions* gpu_options = config->mutable_gpu_options();
  auto virtual_devices =
      gpu_options->mutable_experimental()->add_virtual_devices();
  virtual_devices->add_memory_limit_mb(200);
  virtual_devices->add_memory_limit_mb(200);
  std::unique_ptr<Session> session(NewSession(options));

  {
    // params.cluster is not set, should find and return first gpu id and
    // corresponding allocator.
    auto result = GetDeviceAndAllocator(params, engine_info);
    EXPECT_EQ(0, result.first);
    EXPECT_NE(nullptr, result.second);
    EXPECT_EQ("GPU_0_bfc", result.second->Name());
  }

  FakeCluster cluster;
  params.cluster = &cluster;
  {
    // params.cluster->GetDeviceSet() returns null, should find and return first
    // gpu id and corresponding allocator.
    auto result = GetDeviceAndAllocator(params, engine_info);
    EXPECT_EQ(0, result.first);
    EXPECT_NE(nullptr, result.second);
    EXPECT_EQ("GPU_0_bfc", result.second->Name());
  }

  // Build the DeviceSet.
  DeviceSet device_set;
  const DeviceMgr* device_mgr = nullptr;
  TF_ASSERT_OK(session->LocalDeviceManager(&device_mgr));
  for (auto d : device_mgr->ListDevices()) {
    device_set.AddDevice(d);
  }
  cluster.SetDeviceSet(&device_set);
  {
    // engine_info.device is not set, should find and return first gpu id and
    // corresponding allocator.
    auto result = GetDeviceAndAllocator(params, engine_info);
    EXPECT_EQ(0, result.first);
    EXPECT_NE(nullptr, result.second);
    EXPECT_EQ("GPU_0_bfc", result.second->Name());
  }

  engine_info.device = "/GPU:1";
  {
    // Set to use second device.
    auto result = GetDeviceAndAllocator(params, engine_info);
    EXPECT_EQ(0, result.first);
    EXPECT_NE(nullptr, result.second);
    EXPECT_EQ("GPU_1_bfc", result.second->Name());
  }

  engine_info.device = "/GPU:3";
  {
    // Set to use nonexistent device.
    auto result = GetDeviceAndAllocator(params, engine_info);
    EXPECT_EQ(-1, result.first);
    EXPECT_EQ(nullptr, result.second);
  }
}

class ConvertAfterShapesTest : public ::testing::Test {
 public:
  Status RunConvertAfterShape(Scope s, GraphDef* output_graph_def) {
    // Create GraphProperties.
    grappler::GrapplerItem item;
    TF_EXPECT_OK(s.ToGraphDef(&item.graph));
    grappler::GraphProperties graph_properties(item);
    TF_EXPECT_OK(graph_properties.InferStatically(true));

    // Construct ConversionParams.
    const std::vector<string> output_names{"output"};
    ConversionParams params;
    params.output_names = &output_names;
    params.max_workspace_size_bytes = 8 << 20;
    params.output_graph_def = output_graph_def;
    params.minimum_segment_size = 1;
    params.grappler_item = &item;
    params.use_calibration = false;
    params.trt_logger_name = "DefaultLogger";

    return ConvertAfterShapes(params);
  }
};

TEST_F(ConvertAfterShapesTest, DirectlyConnectedEngines) {
  // Create the graph. There will be two TRTEngineOps after the conversion, and
  // the upstream TRTEngineOp will have two output connections from the same
  // node:port inside the op to the downstream TRTEngineOp. Then, if it adds the
  // downstream TRTEngineOp first, when adding the upstream op it'll need to
  // update the same output connection twice. This test ensures the correctness
  // of the conversion under such condition.
  Scope s = Scope::NewRootScope();
  auto input = ops::Placeholder(s.WithOpName("input"), DT_FLOAT,
                                ops::Placeholder::Shape({2, 1}));
  // We purposefully choose the name of the root node of each segment, so it'll
  // process the segment in the downstream first, then, when it tries to update
  // the edge between the two TRTEngineOps, it'll try to add the same edge
  // multiple times.
  auto segment_root_1 = ops::Identity(s.WithOpName("segment_root_b"), input);
  auto add1 = ops::Add(s.WithOpName("add1"), segment_root_1, segment_root_1);
  // Add incompatible reshapes that change the batch dimension.
  auto incompatible =
      ops::Reshape(s.WithOpName("reshape1"), add1, Input({1, 2}));
  incompatible =
      ops::Reshape(s.WithOpName("reshape2"), incompatible, Input({2, 1}));

  auto add2 = ops::Add(s.WithOpName("add2"), incompatible, add1);
  auto segment_root_2 = ops::Identity(s.WithOpName("segment_root_a"), add1);
  auto add3 = ops::Add(s.WithOpName("add3"), add2, segment_root_2);
  ops::Identity(s.WithOpName("output"), add3);

  GraphDef output_graph_def;
  TF_EXPECT_OK(RunConvertAfterShape(s, &output_graph_def));

  auto remove_graph_sequence_number = [](std::string node_name) {
    const std::regex pattern("TRTEngineOp_[0-9]+_");
    return std::regex_replace(node_name, pattern, "TRTEngineOp_");
  };
  int num_trt_ops = 0;
  for (const NodeDef& node : output_graph_def.node()) {
    std::string node_name = node.name();
    if (node.op() != "TRTEngineOp") continue;
    node_name = remove_graph_sequence_number(node_name);
    if (node_name == "TRTEngineOp_1") {
      EXPECT_EQ(1, node.input_size());
      EXPECT_EQ("input", node.input(0));
      ++num_trt_ops;
    } else if (node_name == "TRTEngineOp_0") {
      EXPECT_EQ(2, node.input_size());
      EXPECT_EQ("TRTEngineOp_1", remove_graph_sequence_number(node.input(0)));
      EXPECT_EQ("reshape2", node.input(1));
      ++num_trt_ops;
    }
  }
  EXPECT_EQ(2, num_trt_ops);
}

}  // namespace convert
}  // namespace tensorrt
}  // namespace tensorflow

#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
<|endoftext|>


<commit_before>/* Copyright 2016 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/distributed_runtime/rpc/grpc_tensor_coding.h"
#include "grpc++/support/byte_buffer.h"
#include "grpc++/support/slice.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/framework/tensor_reference.h"
#include "tensorflow/core/lib/gtl/inlined_vector.h"
#include "tensorflow/core/lib/io/proto_encode_helper.h"
#include "tensorflow/core/platform/env.h"
#include "tensorflow/core/protobuf/worker.pb.h"

namespace tensorflow {
namespace grpc {

static void do_nothing(void* raw) {}
static void unref_tensorreference(void* raw) {
  TensorReference* ref = static_cast<TensorReference*>(raw);
  ref->Unref();
  delete ref;
}

void EncodeRecvTensorResponseToByteBuffer(const RecvTensorResponse& proto,
                                          ::grpc::ByteBuffer* result) {
  size_t len = proto.ByteSize();
  gpr_slice s = gpr_slice_malloc(len);
  proto.SerializeWithCachedSizesToArray(
      reinterpret_cast<uint8*>(GPR_SLICE_START_PTR(s)));
  ::grpc::Slice slice(s, ::grpc::Slice::STEAL_REF);
  *result = ::grpc::ByteBuffer(&slice, 1);
}

// We generate a RecvTensorResponse protocol buffer encoding into "*result",
// but where possible, we share the underlying Tensor buffer for "val", to
// avoid an extra copy.
//
// We hand-encode the protocol buffer data in the following order, as follows:
//
// Let R be a RecvTensorResponse object we want to encode, logically
// constructed by filling in data from "is_dead" and "val" and filling
// in a few other fields as well.
//
// (Letters here are used in the code to refer back to which part of the
//  encoding the code is generating).
//
// A:   <protocol buffer encoding of fields except R.tensor()>
// B1:  <tag encoding for RecvTensorResponse::tensor>
// B2:  <varint32 length of R.tensor() sub message>
// C:   <protocol buffer encoding of R.tensor() except for
//          R.tensor().tensor_content()>
// D1:  <tag encoding for TensorProto::tensor_content>
// D2:  <varint32 length of R.tensor().tensor_content() data>
// E:   <actual data for val's representation>
//
// If the tensor data is up to "kLargeTensorBytes", then A
// through E will all be encoded into "*result" in a single gpr_slice.
//
// If the tensor data is larger than "kLargeTensorBytes", then A through
// D2 will be encoded in one gpr_slice, and E will be encoded in a second
// gpr_slice that points to the backing store for the tensor data, to avoid
// copying the tensor data (and the gpr_slice setup will be arrange so as
// to dereference the underlying tensor data buffer when it is no longer
// needed in the "*result" ByteBuffer).
static int VarLengthEncodingSize(uint32 tag, size_t bytes) {
  return core::VarintLength(tag << 3) + core::VarintLength(bytes) + bytes;
}

void EncodeTensorToByteBuffer(bool is_dead, const Tensor& val,
                              ::grpc::ByteBuffer* result) {
  const int kLargeTensorBytes = 1024;
  RecvTensorResponse response;
  if (is_dead) {
    response.set_is_dead(is_dead);
  }
  response.set_send_start_micros(Env::Default()->NowMicros());
  if (!DataTypeCanUseMemcpy(val.dtype())) {
    // Straightforward but slow path for complicated kinds of tensor data
    // TODO(jeff,sanjay): If this becomes an issue, we could
    // go directly from val -> ByteBuffer, with some effort.
    val.AsProtoTensorContent(response.mutable_tensor());

    // Encode full protocol buffer to a ByteBuffer
    EncodeRecvTensorResponseToByteBuffer(response, result);
  } else {
    TensorProto skeleton;
    skeleton.set_dtype(val.dtype());
    // TODO(jeff): Encode val.shape() directly to protocol buffer rather than
    // going through TensorShapeProto.
    val.shape().AsProto(skeleton.mutable_tensor_shape());

    StringPiece tdata = val.tensor_data();
    uint32 overall_tensor_proto_bytesize =
        (skeleton.ByteSize() +
         VarLengthEncodingSize(TensorProto::kTensorContentFieldNumber,
                               tdata.size()));

    string header;  // All of RecvTensorRequest except the tensor() field
    response.AppendToString(&header);

    string tensor_except_contents;  // tensor() field except contents
    skeleton.AppendToString(&tensor_except_contents);

    size_t expected_size =
        (header.size() +
         VarLengthEncodingSize(RecvTensorResponse::kTensorFieldNumber,
                               overall_tensor_proto_bytesize));
    // If "tensor_data_is_large == false", we copy the tensor data to the
    // end of the buffer we are preparing that holds the rest of the
    // RecvTensorResponse protocol buffer.
    //
    // If "tensor_data_is_large == true", we arrange to share the backing
    // store of the data by creating a slice that also points to the
    // backing store, with appropriate reference counts to keep the
    // backing store alive as needed.
    bool tensor_data_is_large = (tdata.size() > kLargeTensorBytes);
    size_t encoder_size = expected_size - tdata.size();

    // Encode all but the actual "tdata", but including the tag and
    // varlength header for the "tdata"
    gtl::InlinedVector<char, 1024> space(encoder_size);
    io::ProtoEncodeHelper e(space.data(), space.size());
    // (A)
    e.WriteRawBytes(header);

    // (B1) & (B2)
    e.WriteVarlengthBeginning(RecvTensorResponse::kTensorFieldNumber,
                              overall_tensor_proto_bytesize);
    // (C)
    e.WriteRawBytes(tensor_except_contents);
    // (D1) & (D2)
    e.WriteVarlengthBeginning(TensorProto::kTensorContentFieldNumber,
                              tdata.size());

    // All but the tensor backing store are serialized now

    // Now allocate memory and put into the ByteBuffer
    ::grpc::Slice slices[3];
    int num_slices = 0;
    {
      size_t slice_len = e.size() + (tensor_data_is_large ? 0 : tdata.size());
      gpr_slice s0 = gpr_slice_malloc(slice_len);
      memcpy(GPR_SLICE_START_PTR(s0), e.data(), e.size());
      if (!tensor_data_is_large) {
        // (E)
        memcpy(GPR_SLICE_START_PTR(s0) + e.size(), tdata.data(), tdata.size());
      }
      slices[0] = ::grpc::Slice(s0, ::grpc::Slice::STEAL_REF);
      num_slices += 1;
    }

    if (tensor_data_is_large) {
      // Encode the actual tensor data by pointing to the backing store,
      // and add a special zero-length slice that is really a TensorReference
      // object that we will destroy when we are done.
      //
      // TODO(jeff): Note that this approach relies on the fact that
      // slices are destroyed in the order in which they are added to
      // the ByteBuffer.  In principle, these could be broken by future
      // hypothetical grpc_slice-related changes (e.g. the
      // implementation could decide to destroy 0-length slices
      // eagerly).  In practice, this does not happen with the current
      // implementation, and the gpr_slice interface at the moment does
      // not allow us to do the Tensor-unreferencing in the right way
      // (since the Tensor pointer is different than the backing store
      // array pointer).
      //
      // TODO(jeff,sanjay): switch to using new
      // gsr_slice_new_with_user_data interface that allows for
      // different pointers for the data and the argument to the
      // destroy function once that has been added integrated into grpc
      // (see https://github.com/grpc/grpc/pull/7488)

      // (E) Encode tensor data, but by sharing backing store

      // TODO(jeff,sanjay): It'd be nice to avoid this TensorReference
      // allocation, and instead get our hands on the underlying
      // TensorBuffer object and just directly ref it here and unref
      // it in unref_tensorreference.
      TensorReference* ref = new TensorReference(val);
      gpr_slice s1 = gpr_slice_new(
          const_cast<void*>(static_cast<const void*>(tdata.data())),
          tdata.size(), do_nothing);
      slices[1] = ::grpc::Slice(s1, ::grpc::Slice::STEAL_REF);

      gpr_slice s2 = gpr_slice_new(ref, 0, unref_tensorreference);
      slices[2] = ::grpc::Slice(s2, ::grpc::Slice::STEAL_REF);
      num_slices += 2;
    }
    size_t total_bytes = 0;
    for (int i = 0; i < num_slices; i++) {
      total_bytes += slices[i].size();
    }
    CHECK_EQ(total_bytes, expected_size);

    *result = ::grpc::ByteBuffer(&slices[0], num_slices);
  }
}

}  // namespace grpc
}  // namespace tensorflow
<commit_msg>Reduce number of allocations in EncodeTensorToByteBuffer by hand-encoding the skeleton portion of the TensorProto (in particular, encoding the TensorShapeProto was allocating one object per dimension plus an allocation per Tensor).<commit_after>/* Copyright 2016 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/distributed_runtime/rpc/grpc_tensor_coding.h"
#include "grpc++/support/byte_buffer.h"
#include "grpc++/support/slice.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/framework/tensor_reference.h"
#include "tensorflow/core/lib/gtl/inlined_vector.h"
#include "tensorflow/core/lib/io/proto_encode_helper.h"
#include "tensorflow/core/platform/env.h"
#include "tensorflow/core/protobuf/worker.pb.h"

namespace tensorflow {
namespace grpc {

static void do_nothing(void* raw) {}
static void unref_tensorreference(void* raw) {
  TensorReference* ref = static_cast<TensorReference*>(raw);
  ref->Unref();
  delete ref;
}

void EncodeRecvTensorResponseToByteBuffer(const RecvTensorResponse& proto,
                                          ::grpc::ByteBuffer* result) {
  size_t len = proto.ByteSize();
  gpr_slice s = gpr_slice_malloc(len);
  proto.SerializeWithCachedSizesToArray(
      reinterpret_cast<uint8*>(GPR_SLICE_START_PTR(s)));
  ::grpc::Slice slice(s, ::grpc::Slice::STEAL_REF);
  *result = ::grpc::ByteBuffer(&slice, 1);
}

// We generate a RecvTensorResponse protocol buffer encoding into "*result",
// but where possible, we share the underlying Tensor buffer for "val", to
// avoid an extra copy.
//
// We hand-encode the protocol buffer data in the following order, as follows:
//
// Let R be a RecvTensorResponse object we want to encode, logically
// constructed by filling in data from "is_dead" and "val" and filling
// in a few other fields as well.
//
// (Letters here are used in the code to refer back to which part of the
//  encoding the code is generating).
//
// A:   <protocol buffer encoding of fields except R.tensor()>
// B1:  <tag encoding for RecvTensorResponse::tensor>
// B2:  <varint32 length of R.tensor() sub message>
// C:   <protocol buffer encoding of R.tensor() except for
//          R.tensor().tensor_content()>
// D1:  <tag encoding for TensorProto::tensor_content>
// D2:  <varint32 length of R.tensor().tensor_content() data>
// E:   <actual data for val's representation>
//
// If the tensor data is up to "kLargeTensorBytes", then A
// through E will all be encoded into "*result" in a single gpr_slice.
//
// If the tensor data is larger than "kLargeTensorBytes", then A through
// D2 will be encoded in one gpr_slice, and E will be encoded in a second
// gpr_slice that points to the backing store for the tensor data, to avoid
// copying the tensor data (and the gpr_slice setup will be arrange so as
// to dereference the underlying tensor data buffer when it is no longer
// needed in the "*result" ByteBuffer).
static int VarLengthEncodingSize(uint32 tag, size_t bytes) {
  return core::VarintLength(tag << 3) + core::VarintLength(bytes) + bytes;
}

// Returns an upper bound in bytes of the protocol buffer encoding of
// the "skeleton" of "val" (all the data needed for dtype and the shape,
// but not the actual contents of "val").
static int SkeletonEncodingSizeUpperBound(const Tensor& val) {
  static const int kVarintMax64 = 10;  // Max length of varint64 encoding
  const int ndims = val.shape().dims();
  return (2 * kVarintMax64) +           // dtype
         (ndims * (4 * kVarintMax64));  // Shape: 4 varints per dim
}

// Encode the skeleton for "val" (the encoded TensorProto contents
// (dtype and shape, but not the actual data) into "*e".  The backing
// store for "*e" must be of appropriate size to hold this encoding.
static void EncodeSkeleton(const Tensor& val, io::ProtoEncodeHelper* e) {
  // Encode val.dtype()
  e->WriteUint64(TensorProto::kDtypeFieldNumber, val.dtype());

  // Compute length of val.shape() proto encoding
  const int ndims = val.shape().dims();
  int tensor_shape_bytes = 0;
  for (int d = 0; d < ndims; d++) {
    int64 dim_size = val.shape().dim_size(d);
    tensor_shape_bytes +=
        2 +  // TensorShapeProto dim tag + varintlength of submessage
        1 +  // TensorShapeProto_Dim::kSizeFieldNumber
        core::VarintLength(dim_size);
  }

  if (tensor_shape_bytes > 0) {
    e->WriteVarlengthBeginning(TensorProto::kTensorShapeFieldNumber,
                               tensor_shape_bytes);
    // Encode val.shape()
    for (int d = 0; d < ndims; d++) {
      int64 dim_size = val.shape().dim_size(d);
      int64 dim_varlen = 1 +  // TensorShapeProto_Dim::kSizeFieldNumber
                         core::VarintLength(dim_size);
      e->WriteVarlengthBeginning(TensorShapeProto::kDimFieldNumber, dim_varlen);
      e->WriteUint64(TensorShapeProto_Dim::kSizeFieldNumber, dim_size);
    }
  }

#ifndef NDEBUG
  {
    // Debug-mode only check to make sure the encoding above is
    // identical to the auto-generated protocol buffer encoding.
    TensorProto skeleton;
    skeleton.set_dtype(val.dtype());
    val.shape().AsProto(skeleton.mutable_tensor_shape());
    string tensor_except_contents;  // tensor() field except contents
    skeleton.AppendToString(&tensor_except_contents);
    TensorProto skeleton2;
    skeleton2.ParseFromString(string(e->data(), e->size()));
    string out;
    skeleton.AppendToString(&out);
    DCHECK_EQ(tensor_except_contents, out) << skeleton.DebugString() << " vs\n"
                                           << skeleton2.DebugString();
  }
#endif
}

void EncodeTensorToByteBuffer(bool is_dead, const Tensor& val,
                              ::grpc::ByteBuffer* result) {
  const int kLargeTensorBytes = 1024;
  RecvTensorResponse response;
  if (is_dead) {
    response.set_is_dead(is_dead);
  }
  response.set_send_start_micros(Env::Default()->NowMicros());
  if (!DataTypeCanUseMemcpy(val.dtype())) {
    // Straightforward but slow path for complicated kinds of tensor data
    // TODO(jeff,sanjay): If this becomes an issue, we could
    // go directly from val -> ByteBuffer, with some effort.
    val.AsProtoTensorContent(response.mutable_tensor());

    // Encode full protocol buffer to a ByteBuffer
    EncodeRecvTensorResponseToByteBuffer(response, result);
  } else {
    // skeleton is the encoded TensorProto contents (dtype and shape), but
    // not the actual data
    gtl::InlinedVector<char, 128> skeleton(SkeletonEncodingSizeUpperBound(val));
    io::ProtoEncodeHelper e_skeleton(skeleton.data(), skeleton.size());
    EncodeSkeleton(val, &e_skeleton);

    StringPiece tdata = val.tensor_data();
    uint32 overall_tensor_proto_bytesize =
        (e_skeleton.size() +
         VarLengthEncodingSize(TensorProto::kTensorContentFieldNumber,
                               tdata.size()));
    string header;  // All of RecvTensorRequest except the tensor() field
    response.AppendToString(&header);

    size_t expected_size =
        (header.size() +
         VarLengthEncodingSize(RecvTensorResponse::kTensorFieldNumber,
                               overall_tensor_proto_bytesize));
    // If "tensor_data_is_large == false", we copy the tensor data to the
    // end of the buffer we are preparing that holds the rest of the
    // RecvTensorResponse protocol buffer.
    //
    // If "tensor_data_is_large == true", we arrange to share the backing
    // store of the data by creating a slice that also points to the
    // backing store, with appropriate reference counts to keep the
    // backing store alive as needed.
    bool tensor_data_is_large = (tdata.size() > kLargeTensorBytes);
    size_t encoder_size = expected_size - tdata.size();

    // Encode all but the actual "tdata", but including the tag and
    // varlength header for the "tdata"
    gtl::InlinedVector<char, 1024> space(encoder_size);
    io::ProtoEncodeHelper e(space.data(), space.size());
    // (A)
    e.WriteRawBytes(header);

    // (B1) & (B2)
    e.WriteVarlengthBeginning(RecvTensorResponse::kTensorFieldNumber,
                              overall_tensor_proto_bytesize);
    // (C)
    e.WriteRawBytes(StringPiece(e_skeleton.data(), e_skeleton.size()));
    // (D1) & (D2)
    e.WriteVarlengthBeginning(TensorProto::kTensorContentFieldNumber,
                              tdata.size());

    // All but the tensor backing store are serialized now

    // Now allocate memory and put into the ByteBuffer
    ::grpc::Slice slices[3];
    int num_slices = 0;
    {
      size_t slice_len = e.size() + (tensor_data_is_large ? 0 : tdata.size());
      gpr_slice s0 = gpr_slice_malloc(slice_len);
      memcpy(GPR_SLICE_START_PTR(s0), e.data(), e.size());
      if (!tensor_data_is_large) {
        // (E)
        memcpy(GPR_SLICE_START_PTR(s0) + e.size(), tdata.data(), tdata.size());
      }
      slices[0] = ::grpc::Slice(s0, ::grpc::Slice::STEAL_REF);
      num_slices += 1;
    }

    if (tensor_data_is_large) {
      // Encode the actual tensor data by pointing to the backing store,
      // and add a special zero-length slice that is really a TensorReference
      // object that we will destroy when we are done.
      //
      // TODO(jeff): Note that this approach relies on the fact that
      // slices are destroyed in the order in which they are added to
      // the ByteBuffer.  In principle, these could be broken by future
      // hypothetical grpc_slice-related changes (e.g. the
      // implementation could decide to destroy 0-length slices
      // eagerly).  In practice, this does not happen with the current
      // implementation, and the gpr_slice interface at the moment does
      // not allow us to do the Tensor-unreferencing in the right way
      // (since the Tensor pointer is different than the backing store
      // array pointer).
      //
      // TODO(jeff,sanjay): switch to using new
      // gsr_slice_new_with_user_data interface that allows for
      // different pointers for the data and the argument to the
      // destroy function once that has been added integrated into grpc
      // (see https://github.com/grpc/grpc/pull/7488)

      // (E) Encode tensor data, but by sharing backing store

      // TODO(jeff,sanjay): It'd be nice to avoid this TensorReference
      // allocation, and instead get our hands on the underlying
      // TensorBuffer object and just directly ref it here and unref
      // it in unref_tensorreference.
      TensorReference* ref = new TensorReference(val);
      gpr_slice s1 = gpr_slice_new(
          const_cast<void*>(static_cast<const void*>(tdata.data())),
          tdata.size(), do_nothing);
      slices[1] = ::grpc::Slice(s1, ::grpc::Slice::STEAL_REF);

      gpr_slice s2 = gpr_slice_new(ref, 0, unref_tensorreference);
      slices[2] = ::grpc::Slice(s2, ::grpc::Slice::STEAL_REF);
      num_slices += 2;
    }
    size_t total_bytes = 0;
    for (int i = 0; i < num_slices; i++) {
      total_bytes += slices[i].size();
    }
    CHECK_EQ(total_bytes, expected_size);

    *result = ::grpc::ByteBuffer(&slices[0], num_slices);
  }
}

}  // namespace grpc
}  // namespace tensorflow
<|endoftext|>


<commit_before>#include <cassert>
#include <vector>
#include <cmath>

#include "Encoder_RSC_generic_json_sys.hpp"

template <typename B>
Encoder_RSC_generic_json_sys<B>
::Encoder_RSC_generic_json_sys(const int& K, const int& N, const int& n_frames, const bool buffered_encoding,
                               const std::vector<int> poly, std::ostream &stream, const std::string name)
: Encoder_RSC_generic_sys<B>(K, N, n_frames, buffered_encoding, poly, name),
  stream(stream), bit_counter(0), natural_domain(true), poly(poly)
{
}

template <typename B>
std::vector<std::vector<int>> Encoder_RSC_generic_json_sys<B>
::get_trellis()
{
	std::vector<std::vector<int>> trellis(10, std::vector<int>(this->n_states));

	std::vector<bool> occurrence(this->n_states, false);

	for (auto i = 0; i < this->n_states; i++)
	{
		// sys = 0
		auto state   = i;
		auto bit_sys = 0;
		auto bit_par = Encoder_RSC_generic_sys<B>::inner_encode(bit_sys, state);

		trellis[0 + (occurrence[state] ? 3 : 0)][state] = i;                 // initial state
		trellis[1 + (occurrence[state] ? 3 : 0)][state] = +1;                // gamma coeff
		trellis[2 + (occurrence[state] ? 3 : 0)][state] = bit_sys ^ bit_par; // gamma
		trellis[6                              ][i    ] = state;             // final state, bit syst = 0
		trellis[7                              ][i    ] = bit_sys ^ bit_par; // gamma      , bit syst = 1

		occurrence[state] = true;

		// sys = 1
		state   = i;
		bit_sys = 1;
		bit_par = Encoder_RSC_generic_sys<B>::inner_encode(bit_sys, state);

		trellis[0 + (occurrence[state] ? 3 : 0)][state] = i;                 // initial state
		trellis[1 + (occurrence[state] ? 3 : 0)][state] = -1;                // gamma coeff
		trellis[2 + (occurrence[state] ? 3 : 0)][state] = bit_sys ^ bit_par; // gamma
		trellis[8                              ][i    ] = state;             // initial state, bit syst = 0
		trellis[9                              ][i    ] = bit_sys ^ bit_par; // gamma        , bit syst = 1

		occurrence[state] = true;
	}

	return trellis;
}

template <typename B>
int Encoder_RSC_generic_json_sys<B>
::inner_encode(const int bit_sys, int &state)
{
	const auto bit_par = Encoder_RSC_generic_sys<B>::inner_encode(bit_sys, state);

	if (bit_counter == 0)
	{
		if (natural_domain)
		{
			assert(this->n_frames == 1);
			const auto trellis = this->get_trellis();

			stream << "\t" << "[{" << std::endl;

			stream << "\t\t" << "\"stage\": \"encoder\"," << std::endl
			       << "\t\t" << "\"K\": " << this->K << "," << std::endl
			       << "\t\t" << "\"poly\": \"{0" << std::oct << poly[0] << ",0" << std::oct << poly[1] << "}\"," << std::endl
			       << "\t\t" << "\"R\": \"1/3\", " << std::endl
			       << "\t\t" << "\"ff\": " << std::dec << this->n_ff << "," << std::endl
			       << "\t\t" << "\"states\": " << this->n_states << "," << std::endl;

			stream << "\t\t" << "\"trellis\": [" << std::endl;
			for (auto i = 0; i < (int)trellis.size(); i++)
			{
				stream << "\t\t\t" << "[";
				for (auto j = 0; j < (int)trellis[i].size() -1; j++)
					stream << std::setw(3) << trellis[i][j] << ",";
				stream << std::setw(3) << trellis[i][(int)trellis[i].size() -1] << "]";
				if (i != (int)trellis.size() -1)
					stream << "," << std::endl;
			}
			stream << "]," << std::endl;

			stream << "\t\t" << "\"transitions\": {" << std::endl
			       << "\t\t\t" << "\"natural\": [" << std::endl;
		}
		else
			stream << "\t\t\t" << "\"interleaved\": [" << std::endl;
	}

	stream << "\t\t\t\t" << std::setw(2) << state;

	if (bit_counter != (this->K + this->n_ff -1))
	{
		stream << "," << std::endl;
		bit_counter++;
	}
	else
	{
		if (natural_domain)
			stream << "]," << std::endl;
		else
			stream << "]}"  << std::endl
                   << "\t" << "},{" << std::endl;

		bit_counter = 0;
		natural_domain = !natural_domain;
	}

	return bit_par;
}

// ==================================================================================== explicit template instantiation
#include "Tools/types.h"
#ifdef MULTI_PREC
template class Encoder_RSC_generic_json_sys<B_8>;
template class Encoder_RSC_generic_json_sys<B_16>;
template class Encoder_RSC_generic_json_sys<B_32>;
template class Encoder_RSC_generic_json_sys<B_64>;
#else
template class Encoder_RSC_generic_json_sys<B>;
#endif
// ==================================================================================== explicit template instantiation
<commit_msg>Cosmetics.<commit_after>#include <cassert>
#include <vector>
#include <cmath>

#include "Encoder_RSC_generic_json_sys.hpp"

template <typename B>
Encoder_RSC_generic_json_sys<B>
::Encoder_RSC_generic_json_sys(const int& K, const int& N, const int& n_frames, const bool buffered_encoding,
                               const std::vector<int> poly, std::ostream &stream, const std::string name)
: Encoder_RSC_generic_sys<B>(K, N, n_frames, buffered_encoding, poly, name),
  stream(stream), bit_counter(0), natural_domain(true), poly(poly)
{
}

template <typename B>
std::vector<std::vector<int>> Encoder_RSC_generic_json_sys<B>
::get_trellis()
{
	std::vector<std::vector<int>> trellis(10, std::vector<int>(this->n_states));

	std::vector<bool> occurrence(this->n_states, false);

	for (auto i = 0; i < this->n_states; i++)
	{
		// sys = 0
		auto state   = i;
		auto bit_sys = 0;
		auto bit_par = Encoder_RSC_generic_sys<B>::inner_encode(bit_sys, state);

		trellis[0 + (occurrence[state] ? 3 : 0)][state] = i;                 // initial state
		trellis[1 + (occurrence[state] ? 3 : 0)][state] = +1;                // gamma coeff
		trellis[2 + (occurrence[state] ? 3 : 0)][state] = bit_sys ^ bit_par; // gamma
		trellis[6                              ][i    ] = state;             // final state, bit syst = 0
		trellis[7                              ][i    ] = bit_sys ^ bit_par; // gamma      , bit syst = 0

		occurrence[state] = true;

		// sys = 1
		state   = i;
		bit_sys = 1;
		bit_par = Encoder_RSC_generic_sys<B>::inner_encode(bit_sys, state);

		trellis[0 + (occurrence[state] ? 3 : 0)][state] = i;                 // initial state
		trellis[1 + (occurrence[state] ? 3 : 0)][state] = -1;                // gamma coeff
		trellis[2 + (occurrence[state] ? 3 : 0)][state] = bit_sys ^ bit_par; // gamma
		trellis[8                              ][i    ] = state;             // initial state, bit syst = 1
		trellis[9                              ][i    ] = bit_sys ^ bit_par; // gamma        , bit syst = 1

		occurrence[state] = true;
	}

	return trellis;
}

template <typename B>
int Encoder_RSC_generic_json_sys<B>
::inner_encode(const int bit_sys, int &state)
{
	const auto bit_par = Encoder_RSC_generic_sys<B>::inner_encode(bit_sys, state);

	if (bit_counter == 0)
	{
		if (natural_domain)
		{
			assert(this->n_frames == 1);
			const auto trellis = this->get_trellis();

			stream << "\t" << "[{" << std::endl;

			stream << "\t\t" << "\"stage\": \"encoder\"," << std::endl
			       << "\t\t" << "\"K\": " << this->K << "," << std::endl
			       << "\t\t" << "\"poly\": \"{0" << std::oct << poly[0] << ",0" << std::oct << poly[1] << "}\"," << std::endl
			       << "\t\t" << "\"R\": \"1/3\", " << std::endl
			       << "\t\t" << "\"ff\": " << std::dec << this->n_ff << "," << std::endl
			       << "\t\t" << "\"states\": " << this->n_states << "," << std::endl;

			stream << "\t\t" << "\"trellis\": [" << std::endl;
			for (auto i = 0; i < (int)trellis.size(); i++)
			{
				stream << "\t\t\t" << "[";
				for (auto j = 0; j < (int)trellis[i].size() -1; j++)
					stream << std::setw(3) << trellis[i][j] << ",";
				stream << std::setw(3) << trellis[i][(int)trellis[i].size() -1] << "]";
				if (i != (int)trellis.size() -1)
					stream << "," << std::endl;
			}
			stream << "]," << std::endl;

			stream << "\t\t" << "\"transitions\": {" << std::endl
			       << "\t\t\t" << "\"natural\": [" << std::endl;
		}
		else
			stream << "\t\t\t" << "\"interleaved\": [" << std::endl;
	}

	stream << "\t\t\t\t" << std::setw(2) << state;

	if (bit_counter != (this->K + this->n_ff -1))
	{
		stream << "," << std::endl;
		bit_counter++;
	}
	else
	{
		if (natural_domain)
			stream << "]," << std::endl;
		else
			stream << "]}"  << std::endl
			       << "\t" << "},{" << std::endl;

		bit_counter = 0;
		natural_domain = !natural_domain;
	}

	return bit_par;
}

// ==================================================================================== explicit template instantiation
#include "Tools/types.h"
#ifdef MULTI_PREC
template class Encoder_RSC_generic_json_sys<B_8>;
template class Encoder_RSC_generic_json_sys<B_16>;
template class Encoder_RSC_generic_json_sys<B_32>;
template class Encoder_RSC_generic_json_sys<B_64>;
#else
template class Encoder_RSC_generic_json_sys<B>;
#endif
// ==================================================================================== explicit template instantiation
<|endoftext|>