shed: bitfield intersect command

cmd/lotus-shed/bitfield.go

@@ -28,6 +28,7 @@ var bitFieldCmd = &cli.Command{
 		bitFieldRunsCmd,
 		bitFieldStatCmd,
 		bitFieldDecodeCmd,
+		bitFieldIntersectCmd,
 	},
 }
 
@@ -200,38 +201,9 @@ var bitFieldDecodeCmd = &cli.Command{
 		},
 	},
 	Action: func(cctx *cli.Context) error {
-		var val string
-		if cctx.Args().Present() {
-			val = cctx.Args().Get(0)
-		} else {
-			b, err := ioutil.ReadAll(os.Stdin)
-			if err != nil {
-				return err
-			}
-			val = string(b)
-		}
-
-		var dec []byte
-		switch cctx.String("enc") {
-		case "base64":
-			d, err := base64.StdEncoding.DecodeString(val)
-			if err != nil {
-				return fmt.Errorf("decoding base64 value: %w", err)
-			}
-			dec = d
-		case "hex":
-			d, err := hex.DecodeString(val)
-			if err != nil {
-				return fmt.Errorf("decoding hex value: %w", err)
-			}
-			dec = d
-		default:
-			return fmt.Errorf("unrecognized encoding: %s", cctx.String("enc"))
-		}
-
-		rle, err := bitfield.NewFromBytes(dec)
+		rle, err := decode(cctx, 0)
 		if err != nil {
-			return xerrors.Errorf("failed to parse bitfield: %w", err)
+			return err
 		}
 
 		vals, err := rle.All(100000000000)
@@ -243,3 +215,84 @@ var bitFieldDecodeCmd = &cli.Command{
 		return nil
 	},
 }
+
+var bitFieldIntersectCmd = &cli.Command{
+	Name:        "intersect",
+	Description: "intersect 2 bitfields and print the resulting bitfield as base64",
+	Flags: []cli.Flag{
+		&cli.StringFlag{
+			Name:  "enc",
+			Value: "base64",
+			Usage: "specify input encoding to parse",
+		},
+	},
+	Action: func(cctx *cli.Context) error {
+		b, err := decode(cctx, 1)
+		if err != nil {
+			return err
+		}
+
+		a, err := decode(cctx, 0)
+		if err != nil {
+			return err
+		}
+
+		o, err := bitfield.IntersectBitField(a, b)
+		if err != nil {
+			return xerrors.Errorf("intersect: %w", err)
+		}
+
+		s, err := o.RunIterator()
+		if err != nil {
+			return err
+		}
+
+		bytes, err := rlepluslazy.EncodeRuns(s, []byte{})
+		if err != nil {
+			return err
+		}
+
+		fmt.Println(base64.StdEncoding.EncodeToString(bytes))
+
+		return nil
+	},
+}
+
+func decode(cctx *cli.Context, a int) (bitfield.BitField, error) {
+	var val string
+	if cctx.Args().Present() {
+		if a >= cctx.NArg() {
+			return bitfield.BitField{}, xerrors.Errorf("need more than %d args", a)
+		}
+		val = cctx.Args().Get(a)
+	} else {
+		if a > 0 {
+			return bitfield.BitField{}, xerrors.Errorf("need more than %d args", a)
+		}
+		b, err := ioutil.ReadAll(os.Stdin)
+		if err != nil {
+			return bitfield.BitField{}, err
+		}
+		val = string(b)
+	}
+
+	var dec []byte
+	switch cctx.String("enc") {
+	case "base64":
+		d, err := base64.StdEncoding.DecodeString(val)
+		if err != nil {
+			return bitfield.BitField{}, fmt.Errorf("decoding base64 value: %w", err)
+		}
+		dec = d
+	case "hex":
+		d, err := hex.DecodeString(val)
+		if err != nil {
+			return bitfield.BitField{}, fmt.Errorf("decoding hex value: %w", err)
+		}
+		dec = d
+	default:
+		return bitfield.BitField{}, fmt.Errorf("unrecognized encoding: %s", cctx.String("enc"))
+	}
+
+	return bitfield.NewFromBytes(dec)
+}