snow: yuv444 support

Signed-off-by: Michael Niedermayer <michaelni@gmx.at>

1 files changed, 14 insertions, 11 deletions

diff --git a/libavcodec/snow.h b/libavcodec/snow.h
index c3280154dd..d8a08fa9aa 100644
--- a/libavcodec/snow.h
+++ b/libavcodec/snow.h
@@ -404,20 +404,21 @@ static av_always_inline void predict_slice(SnowContext *s, IDWTELEM *buf, int pl
     const int mb_h= s->b_height << s->block_max_depth;
     int x, y, mb_x;
     int block_size = MB_SIZE >> s->block_max_depth;
-    int block_w    = plane_index ? block_size/2 : block_size;
-    const uint8_t *obmc  = plane_index ? ff_obmc_tab[s->block_max_depth+1] : ff_obmc_tab[s->block_max_depth];
-    const int obmc_stride= plane_index ? block_size : 2*block_size;
+    int block_w    = plane_index ? block_size>>s->chroma_h_shift : block_size;
+    int block_h    = plane_index ? block_size>>s->chroma_v_shift : block_size;
+    const uint8_t *obmc  = plane_index ? ff_obmc_tab[s->block_max_depth+s->chroma_h_shift] : ff_obmc_tab[s->block_max_depth];
+    const int obmc_stride= plane_index ? (2*block_size)>>s->chroma_h_shift : 2*block_size;
     int ref_stride= s->current_picture.linesize[plane_index];
     uint8_t *dst8= s->current_picture.data[plane_index];
     int w= p->width;
     int h= p->height;
-
+    av_assert2(s->chroma_h_shift == s->chroma_v_shift); // obmc params assume squares
     if(s->keyframe || (s->avctx->debug&512)){
         if(mb_y==mb_h)
             return;
 
         if(add){
-            for(y=block_w*mb_y; y<FFMIN(h,block_w*(mb_y+1)); y++){
+            for(y=block_h*mb_y; y<FFMIN(h,block_h*(mb_y+1)); y++){
                 for(x=0; x<w; x++){
                     int v= buf[x + y*w] + (128<<FRAC_BITS) + (1<<(FRAC_BITS-1));
                     v >>= FRAC_BITS;
@@ -426,7 +427,7 @@ static av_always_inline void predict_slice(SnowContext *s, IDWTELEM *buf, int pl
                 }
             }
         }else{
-            for(y=block_w*mb_y; y<FFMIN(h,block_w*(mb_y+1)); y++){
+            for(y=block_h*mb_y; y<FFMIN(h,block_h*(mb_y+1)); y++){
                 for(x=0; x<w; x++){
                     buf[x + y*w]-= 128<<FRAC_BITS;
                 }
@@ -439,8 +440,8 @@ static av_always_inline void predict_slice(SnowContext *s, IDWTELEM *buf, int pl
     for(mb_x=0; mb_x<=mb_w; mb_x++){
         add_yblock(s, 0, NULL, buf, dst8, obmc,
                    block_w*mb_x - block_w/2,
-                   block_w*mb_y - block_w/2,
-                   block_w, block_w,
+                   block_h*mb_y - block_h/2,
+                   block_w, block_h,
                    w, h,
                    w, ref_stride, obmc_stride,
                    mb_x - 1, mb_y - 1,
@@ -460,6 +461,7 @@ static inline void set_blocks(SnowContext *s, int level, int x, int y, int l, in
     const int rem_depth= s->block_max_depth - level;
     const int index= (x + y*w) << rem_depth;
     const int block_w= 1<<rem_depth;
+    const int block_h= 1<<rem_depth; //FIXME "w!=h"
     BlockNode block;
     int i,j;
 
@@ -472,7 +474,7 @@ static inline void set_blocks(SnowContext *s, int level, int x, int y, int l, in
     block.type= type;
     block.level= level;
 
-    for(j=0; j<block_w; j++){
+    for(j=0; j<block_h; j++){
         for(i=0; i<block_w; i++){
             s->block[index + i + j*w]= block;
         }
@@ -480,10 +482,11 @@ static inline void set_blocks(SnowContext *s, int level, int x, int y, int l, in
 }
 
 static inline void init_ref(MotionEstContext *c, uint8_t *src[3], uint8_t *ref[3], uint8_t *ref2[3], int x, int y, int ref_index){
+    SnowContext *s = c->avctx->priv_data;
     const int offset[3]= {
           y*c->  stride + x,
-        ((y*c->uvstride + x)>>1),
-        ((y*c->uvstride + x)>>1),
+        ((y*c->uvstride + x)>>s->chroma_h_shift),
+        ((y*c->uvstride + x)>>s->chroma_h_shift),
     };
     int i;
     for(i=0; i<3; i++){