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 | |
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 | |
35 | void DynamicBuffer::AddString(const StringRef& string) { |
36 | AddString(string.str, string.len); |
37 | } |
38 | |
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 | |
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 (size_t 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 | |
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 |
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 | |
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 | |
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 | |
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 | |
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 | |
142 | StringRef GetString(const TfLiteTensor* tensor, int string_index) { |
143 | return GetString(tensor->data.raw, string_index); |
144 | } |
145 | |
146 | } // namespace tflite |
147 | |