1 /* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
2
3 Licensed under the Apache License, Version 2.0 (the "License");
4 you may not use this file except in compliance with the License.
5 You may obtain a copy of the License at
6
7 http://www.apache.org/licenses/LICENSE-2.0
8
9 Unless required by applicable law or agreed to in writing, software
10 distributed under the License is distributed on an "AS IS" BASIS,
11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 See the License for the specific language governing permissions and
13 limitations under the License.
14 ==============================================================================*/
15
16 #include "tensorflow/lite/string_util.h"
17
18 #include <stddef.h>
19 #include <stdint.h>
20
21 #include <cstdlib>
22 #include <cstring>
23 #include <vector>
24
25 #include "tensorflow/lite/c/common.h"
26
27 namespace tflite {
28
AddString(const char * str,size_t len)29 void DynamicBuffer::AddString(const char* str, size_t len) {
30 data_.resize(data_.size() + len);
31 memcpy(data_.data() + offset_.back(), str, len);
32 offset_.push_back(offset_.back() + len);
33 }
34
AddString(const StringRef & string)35 void DynamicBuffer::AddString(const StringRef& string) {
36 AddString(string.str, string.len);
37 }
38
AddJoinedString(const std::vector<StringRef> & strings,char separator)39 void DynamicBuffer::AddJoinedString(const std::vector<StringRef>& strings,
40 char separator) {
41 StringRef ref;
42 ref.str = &separator;
43 ref.len = 1;
44 AddJoinedString(strings, ref);
45 }
46
AddJoinedString(const std::vector<StringRef> & strings,StringRef separator)47 void DynamicBuffer::AddJoinedString(const std::vector<StringRef>& strings,
48 StringRef separator) {
49 // Resize the data buffer.
50 int total_len = (strings.size() - 1) * separator.len;
51 for (StringRef ref : strings) {
52 total_len += ref.len;
53 }
54 data_.resize(data_.size() + total_len);
55
56 char* dst = data_.data() + offset_.back();
57 for (int i = 0; i < strings.size(); ++i) {
58 // Fill separator if not first string.
59 if (i != 0) {
60 memcpy(dst, separator.str, separator.len);
61 dst += separator.len;
62 }
63
64 // Fill content of the string.
65 memcpy(dst, strings[i].str, strings[i].len);
66 dst += strings[i].len;
67 }
68 offset_.push_back(offset_.back() + total_len);
69 }
70
WriteToBuffer(char ** buffer)71 int DynamicBuffer::WriteToBuffer(char** buffer) {
72 // Allocate sufficient memory to tensor buffer.
73 int32_t num_strings = offset_.size() - 1;
74 // Total bytes include:
75 // * size of content (data_.size)
76 // * offset of each tensor (sizeof(int32_t) * num_strings)
77 // * length of whole buffer (int32_t)
78 // * num of strings (int32_t).
79 int32_t bytes = data_.size() // size of content
80 + sizeof(int32_t) * (num_strings + 2); // size of header
81
82 // Caller will take ownership of buffer.
83 *buffer = reinterpret_cast<char*>(malloc(bytes));
84
85 // Set num of string
86 memcpy(*buffer, &num_strings, sizeof(int32_t));
87
88 // Set offset of strings.
89 int32_t start = sizeof(int32_t) * (num_strings + 2);
90 for (size_t i = 0; i < offset_.size(); i++) {
91 int32_t offset = start + offset_[i];
92 memcpy(*buffer + sizeof(int32_t) * (i + 1), &offset, sizeof(int32_t));
93 }
94
95 // Copy data of strings.
96 memcpy(*buffer + start, data_.data(), data_.size());
97 return bytes;
98 }
99
100 #ifndef TF_LITE_STATIC_MEMORY
WriteToTensorAsVector(TfLiteTensor * tensor)101 void DynamicBuffer::WriteToTensorAsVector(TfLiteTensor* tensor) {
102 auto dims = TfLiteIntArrayCreate(1);
103 dims->data[0] = offset_.size() - 1; // Store number of strings.
104 WriteToTensor(tensor, dims);
105 }
106
WriteToTensor(TfLiteTensor * tensor,TfLiteIntArray * new_shape)107 void DynamicBuffer::WriteToTensor(TfLiteTensor* tensor,
108 TfLiteIntArray* new_shape) {
109 char* tensor_buffer;
110 int bytes = WriteToBuffer(&tensor_buffer);
111
112 if (new_shape == nullptr) {
113 new_shape = TfLiteIntArrayCopy(tensor->dims);
114 }
115
116 // Set tensor content pointer to tensor_buffer, and release original data.
117 TfLiteTensorReset(tensor->type, tensor->name, new_shape, tensor->params,
118 tensor_buffer, bytes, kTfLiteDynamic, tensor->allocation,
119 tensor->is_variable, tensor);
120 }
121 #endif // TF_LITE_STATIC_MEMORY
122
GetStringCount(const void * raw_buffer)123 int GetStringCount(const void* raw_buffer) {
124 // The first integers in the raw buffer is the number of strings.
125 return *static_cast<const int32_t*>(raw_buffer);
126 }
127
GetStringCount(const TfLiteTensor * tensor)128 int GetStringCount(const TfLiteTensor* tensor) {
129 // The first integers in the raw buffer is the number of strings.
130 return GetStringCount(tensor->data.raw);
131 }
132
GetString(const void * raw_buffer,int string_index)133 StringRef GetString(const void* raw_buffer, int string_index) {
134 const int32_t* offset =
135 static_cast<const int32_t*>(raw_buffer) + (string_index + 1);
136 return StringRef{
137 static_cast<const char*>(raw_buffer) + (*offset),
138 (*(offset + 1)) - (*offset),
139 };
140 }
141
GetString(const TfLiteTensor * tensor,int string_index)142 StringRef GetString(const TfLiteTensor* tensor, int string_index) {
143 return GetString(tensor->data.raw, string_index);
144 }
145
146 } // namespace tflite
147