Skip to main content

oximedia_transcode/
segment_encoder.rs

1//! Segment-based encoding for HLS and DASH streaming.
2//!
3//! This module provides tools for planning segment boundaries,
4//! tracking encoded segments, and generating HLS/DASH manifests.
5
6/// Configuration for segment-based encoding.
7#[derive(Debug, Clone)]
8pub struct SegmentConfig {
9    /// Target segment duration in seconds.
10    pub duration_secs: f32,
11    /// Keyframe interval in frames.
12    pub keyframe_interval: u32,
13    /// Force a keyframe at each segment boundary.
14    pub force_key_at_segment: bool,
15}
16
17impl Default for SegmentConfig {
18    fn default() -> Self {
19        Self {
20            duration_secs: 6.0,
21            keyframe_interval: 60,
22            force_key_at_segment: true,
23        }
24    }
25}
26
27impl SegmentConfig {
28    /// Creates a new segment configuration.
29    #[must_use]
30    pub fn new(duration_secs: f32, keyframe_interval: u32, force_key_at_segment: bool) -> Self {
31        Self {
32            duration_secs,
33            keyframe_interval,
34            force_key_at_segment,
35        }
36    }
37}
38
39/// A boundary point between two segments.
40#[derive(Debug, Clone, PartialEq)]
41pub struct SegmentBoundary {
42    /// Frame index where the segment starts.
43    pub frame_idx: u64,
44    /// Whether this frame is a keyframe.
45    pub is_keyframe: bool,
46    /// Timestamp in seconds.
47    pub timestamp_secs: f64,
48}
49
50/// A complete segment plan for encoding.
51#[derive(Debug, Clone)]
52pub struct SegmentPlan {
53    /// All segment boundaries.
54    pub boundaries: Vec<SegmentBoundary>,
55    /// Total number of frames.
56    pub total_frames: u64,
57    /// Total number of segments.
58    pub segment_count: u32,
59}
60
61/// Plans segment boundaries for a given video.
62#[derive(Debug, Clone, Default)]
63pub struct SegmentPlanner;
64
65impl SegmentPlanner {
66    /// Creates a new segment planner.
67    #[must_use]
68    pub fn new() -> Self {
69        Self
70    }
71
72    /// Plans segment boundaries for `total_frames` frames at `fps` with the given config.
73    #[must_use]
74    pub fn plan(total_frames: u64, fps: f32, config: &SegmentConfig) -> SegmentPlan {
75        let frames_per_segment = (config.duration_secs * fps).round() as u64;
76        let frames_per_segment = frames_per_segment.max(1);
77
78        let mut boundaries = Vec::new();
79        let mut frame_idx = 0u64;
80        let mut seg_count = 0u32;
81
82        while frame_idx < total_frames {
83            let is_keyframe = if config.force_key_at_segment {
84                true
85            } else {
86                // Keyframe at regular intervals
87                frame_idx % u64::from(config.keyframe_interval) == 0
88            };
89
90            let timestamp_secs = frame_idx as f64 / f64::from(fps);
91
92            boundaries.push(SegmentBoundary {
93                frame_idx,
94                is_keyframe,
95                timestamp_secs,
96            });
97
98            frame_idx += frames_per_segment;
99            seg_count += 1;
100        }
101
102        SegmentPlan {
103            boundaries,
104            total_frames,
105            segment_count: seg_count,
106        }
107    }
108}
109
110/// An encoded segment of media.
111#[derive(Debug, Clone, PartialEq, Eq)]
112pub struct EncodedSegment {
113    /// Segment index (0-based).
114    pub index: u32,
115    /// Start time in milliseconds.
116    pub start_ms: u64,
117    /// Duration in milliseconds.
118    pub duration_ms: u64,
119    /// File size in bytes.
120    pub size_bytes: u64,
121    /// Actual bitrate in kilobits per second.
122    pub bitrate_kbps: u32,
123    /// Codec used for this segment.
124    pub codec: String,
125}
126
127impl EncodedSegment {
128    /// Creates a new encoded segment.
129    #[must_use]
130    pub fn new(
131        index: u32,
132        start_ms: u64,
133        duration_ms: u64,
134        size_bytes: u64,
135        bitrate_kbps: u32,
136        codec: impl Into<String>,
137    ) -> Self {
138        Self {
139            index,
140            start_ms,
141            duration_ms,
142            size_bytes,
143            bitrate_kbps,
144            codec: codec.into(),
145        }
146    }
147
148    /// Returns the end time in milliseconds.
149    #[must_use]
150    pub fn end_ms(&self) -> u64 {
151        self.start_ms + self.duration_ms
152    }
153
154    /// Returns the duration in seconds.
155    #[must_use]
156    pub fn duration_secs(&self) -> f64 {
157        self.duration_ms as f64 / 1000.0
158    }
159}
160
161/// Encoder that tracks encoded segments.
162#[derive(Debug, Clone, Default)]
163pub struct SegmentEncoder {
164    /// All encoded segments.
165    pub encoded_segments: Vec<EncodedSegment>,
166}
167
168impl SegmentEncoder {
169    /// Creates a new segment encoder.
170    #[must_use]
171    pub fn new() -> Self {
172        Self::default()
173    }
174
175    /// Adds a segment to the encoder's list.
176    pub fn add_segment(&mut self, segment: EncodedSegment) {
177        self.encoded_segments.push(segment);
178    }
179
180    /// Returns the total number of encoded segments.
181    #[must_use]
182    pub fn segment_count(&self) -> usize {
183        self.encoded_segments.len()
184    }
185
186    /// Returns the total encoded size in bytes.
187    #[must_use]
188    pub fn total_bytes(&self) -> u64 {
189        self.encoded_segments.iter().map(|s| s.size_bytes).sum()
190    }
191
192    /// Returns the average bitrate across all segments.
193    #[must_use]
194    pub fn average_bitrate_kbps(&self) -> Option<u32> {
195        if self.encoded_segments.is_empty() {
196            return None;
197        }
198        let sum: u64 = self
199            .encoded_segments
200            .iter()
201            .map(|s| u64::from(s.bitrate_kbps))
202            .sum();
203        Some((sum / self.encoded_segments.len() as u64) as u32)
204    }
205}
206
207/// Generates HLS and DASH manifests from encoded segments.
208#[derive(Debug, Clone, Default)]
209pub struct SegmentManifest;
210
211impl SegmentManifest {
212    /// Generates an HLS `.m3u8` manifest.
213    #[must_use]
214    pub fn generate_hls(segments: &[EncodedSegment], base_url: &str) -> String {
215        let max_duration = segments
216            .iter()
217            .map(EncodedSegment::duration_secs)
218            .fold(0.0_f64, f64::max);
219
220        let mut manifest = format!(
221            "#EXTM3U\n#EXT-X-VERSION:3\n#EXT-X-TARGETDURATION:{}\n#EXT-X-MEDIA-SEQUENCE:0\n",
222            max_duration.ceil() as u64
223        );
224
225        for seg in segments {
226            let duration = seg.duration_secs();
227            manifest.push_str(&format!(
228                "#EXTINF:{:.3},\n{}/segment_{:05}.ts\n",
229                duration, base_url, seg.index
230            ));
231        }
232
233        manifest.push_str("#EXT-X-ENDLIST\n");
234        manifest
235    }
236
237    /// Generates a MPEG-DASH `manifest.mpd` manifest.
238    #[must_use]
239    pub fn generate_dash(segments: &[EncodedSegment], base_url: &str) -> String {
240        let total_ms: u64 = segments.iter().map(|s| s.duration_ms).sum();
241        let total_secs = total_ms as f64 / 1000.0;
242
243        let avg_bitrate = if segments.is_empty() {
244            0u32
245        } else {
246            let sum: u64 = segments.iter().map(|s| u64::from(s.bitrate_kbps)).sum();
247            (sum / segments.len() as u64) as u32
248        };
249
250        let codec = segments.first().map_or("avc1", |s| s.codec.as_str());
251
252        let mut mpd = format!(
253            "<?xml version=\"1.0\" encoding=\"UTF-8\"?>\n\
254             <MPD xmlns=\"urn:mpeg:dash:schema:mpd:2011\" \
255             mediaPresentationDuration=\"PT{total_secs:.3}S\" \
256             type=\"static\">\n  \
257             <Period>\n    \
258             <AdaptationSet mimeType=\"video/mp4\">\n      \
259             <Representation id=\"0\" codecs=\"{codec}\" bandwidth=\"{}\">\n",
260            avg_bitrate * 1000
261        );
262
263        mpd.push_str("        <SegmentList>\n");
264        for seg in segments {
265            mpd.push_str(&format!(
266                "          <SegmentURL media=\"{}/segment_{:05}.mp4\"/>\n",
267                base_url, seg.index
268            ));
269        }
270        mpd.push_str(
271            "        </SegmentList>\n      \
272             </Representation>\n    \
273             </AdaptationSet>\n  \
274             </Period>\n\
275             </MPD>\n",
276        );
277
278        mpd
279    }
280}
281
282#[cfg(test)]
283mod tests {
284    use super::*;
285
286    #[test]
287    fn test_segment_config_default() {
288        let cfg = SegmentConfig::default();
289        assert_eq!(cfg.duration_secs, 6.0);
290        assert!(cfg.force_key_at_segment);
291    }
292
293    #[test]
294    fn test_segment_config_new() {
295        let cfg = SegmentConfig::new(4.0, 120, false);
296        assert_eq!(cfg.duration_secs, 4.0);
297        assert_eq!(cfg.keyframe_interval, 120);
298        assert!(!cfg.force_key_at_segment);
299    }
300
301    #[test]
302    fn test_segment_planner_basic() {
303        let cfg = SegmentConfig::default(); // 6s segments
304                                            // 180 frames at 30fps = 6 seconds total → 1 segment
305        let plan = SegmentPlanner::plan(180, 30.0, &cfg);
306        assert_eq!(plan.total_frames, 180);
307        assert!(plan.segment_count >= 1);
308    }
309
310    #[test]
311    fn test_segment_planner_multiple_segments() {
312        let cfg = SegmentConfig::new(2.0, 30, true);
313        // 60 frames at 30fps = 2s per segment → expect 1 segment start at 0
314        let plan = SegmentPlanner::plan(60, 30.0, &cfg);
315        assert!(!plan.boundaries.is_empty());
316    }
317
318    #[test]
319    fn test_segment_planner_keyframe_at_boundary() {
320        let cfg = SegmentConfig::new(2.0, 60, true);
321        let plan = SegmentPlanner::plan(120, 30.0, &cfg);
322        for b in &plan.boundaries {
323            assert!(b.is_keyframe, "All boundaries should be keyframes");
324        }
325    }
326
327    #[test]
328    fn test_segment_boundary_timestamp() {
329        let cfg = SegmentConfig::new(2.0, 60, true);
330        let plan = SegmentPlanner::plan(120, 30.0, &cfg);
331        assert!((plan.boundaries[0].timestamp_secs - 0.0).abs() < 1e-9);
332    }
333
334    #[test]
335    fn test_encoded_segment_end_ms() {
336        let seg = EncodedSegment::new(0, 0, 2000, 512_000, 2048, "h264");
337        assert_eq!(seg.end_ms(), 2000);
338        assert!((seg.duration_secs() - 2.0).abs() < 1e-9);
339    }
340
341    #[test]
342    fn test_segment_encoder_add_and_count() {
343        let mut enc = SegmentEncoder::new();
344        assert_eq!(enc.segment_count(), 0);
345        enc.add_segment(EncodedSegment::new(0, 0, 2000, 1024, 4000, "h264"));
346        enc.add_segment(EncodedSegment::new(1, 2000, 2000, 2048, 8000, "h264"));
347        assert_eq!(enc.segment_count(), 2);
348    }
349
350    #[test]
351    fn test_segment_encoder_total_bytes() {
352        let mut enc = SegmentEncoder::new();
353        enc.add_segment(EncodedSegment::new(0, 0, 2000, 1000, 4000, "h264"));
354        enc.add_segment(EncodedSegment::new(1, 2000, 2000, 2000, 8000, "h264"));
355        assert_eq!(enc.total_bytes(), 3000);
356    }
357
358    #[test]
359    fn test_segment_encoder_average_bitrate() {
360        let mut enc = SegmentEncoder::new();
361        assert!(enc.average_bitrate_kbps().is_none());
362        enc.add_segment(EncodedSegment::new(0, 0, 2000, 1000, 4000, "h264"));
363        enc.add_segment(EncodedSegment::new(1, 2000, 2000, 2000, 6000, "h264"));
364        assert_eq!(enc.average_bitrate_kbps(), Some(5000));
365    }
366
367    #[test]
368    fn test_generate_hls_contains_extm3u() {
369        let segments = vec![
370            EncodedSegment::new(0, 0, 6000, 1000, 4000, "h264"),
371            EncodedSegment::new(1, 6000, 6000, 1000, 4000, "h264"),
372        ];
373        let manifest = SegmentManifest::generate_hls(&segments, "https://cdn.example.com");
374        assert!(manifest.contains("#EXTM3U"));
375        assert!(manifest.contains("#EXT-X-ENDLIST"));
376        assert!(manifest.contains("segment_00000.ts"));
377        assert!(manifest.contains("segment_00001.ts"));
378    }
379
380    #[test]
381    fn test_generate_dash_contains_mpd() {
382        let segments = vec![EncodedSegment::new(0, 0, 6000, 1000, 4000, "avc1")];
383        let manifest = SegmentManifest::generate_dash(&segments, "https://cdn.example.com");
384        assert!(manifest.contains("<?xml"));
385        assert!(manifest.contains("<MPD"));
386        assert!(manifest.contains("segment_00000.mp4"));
387    }
388
389    #[test]
390    fn test_generate_hls_empty() {
391        let manifest = SegmentManifest::generate_hls(&[], "https://cdn.example.com");
392        assert!(manifest.contains("#EXTM3U"));
393        assert!(manifest.contains("#EXT-X-ENDLIST"));
394    }
395
396    #[test]
397    fn test_generate_dash_empty() {
398        let manifest = SegmentManifest::generate_dash(&[], "https://cdn.example.com");
399        assert!(manifest.contains("<?xml"));
400    }
401}