segment rendering fixes

3 files changed, 112 insertions, 2 deletions

diff --git a/src/display.rs b/src/display.rs
index 50375f2..76db507 100644
--- a/src/display.rs
+++ b/src/display.rs
@@ -10,7 +10,20 @@ use std::hint::unreachable_unchecked;
 use yaxpeax_arch::{Arch, Colorize, ColorSettings, Decodable, LengthedInstruction, ShowContextual, YaxColors};
 use yaxpeax_arch::display::*;
 
-use ::{RegSpec, RegisterBank, Opcode, Operand, Instruction};
+use ::{RegSpec, RegisterBank, Opcode, Operand, Instruction, Segment};
+
+impl fmt::Display for Segment {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        match self {
+            Segment::CS => write!(f, "cs"),
+            Segment::DS => write!(f, "ds"),
+            Segment::ES => write!(f, "es"),
+            Segment::FS => write!(f, "fs"),
+            Segment::GS => write!(f, "gs"),
+            Segment::SS => write!(f, "ss"),
+        }
+    }
+}
 
 impl fmt::Display for RegSpec {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
@@ -537,6 +550,7 @@ impl <T: std::fmt::Write> ShowContextual<u64, [Option<String>], T> for Instructi
             write!(out, "lock ")?;
         }
         self.opcode.colorize(colors, out)?;
+
         match context.and_then(|xs| xs[0].as_ref()) {
             Some(s) => { write!(out, " {}", s)?; },
             None => {
@@ -546,6 +560,9 @@ impl <T: std::fmt::Write> ShowContextual<u64, [Option<String>], T> for Instructi
                     },
                     ref x @ _ => {
                         write!(out, " ")?;
+                        if let Some(prefix) = self.segment_override_for_op(0) {
+                            write!(out, "{}:", prefix);
+                        }
                         x.colorize(colors, out)?;
                     }
                 }
@@ -562,10 +579,16 @@ impl <T: std::fmt::Write> ShowContextual<u64, [Option<String>], T> for Instructi
                             },
                             x @ &Operand::Register(_) => {
                                 write!(out, ", ")?;
+                                if let Some(prefix) = self.segment_override_for_op(1) {
+                                    write!(out, "{}:", prefix);
+                                }
                                 x.colorize(colors, out)
                             }
                             x @ _ => {
                                 write!(out, ", byte ")?;
+                                if let Some(prefix) = self.segment_override_for_op(1) {
+                                    write!(out, "{}:", prefix);
+                                }
                                 x.colorize(colors, out)
                             }
                         }
@@ -582,10 +605,16 @@ impl <T: std::fmt::Write> ShowContextual<u64, [Option<String>], T> for Instructi
                             },
                             x @ &Operand::Register(_) => {
                                 write!(out, ", ")?;
+                                if let Some(prefix) = self.segment_override_for_op(1) {
+                                    write!(out, "{}:", prefix);
+                                }
                                 x.colorize(colors, out)
                             }
                             x @ _ => {
                                 write!(out, ", word ")?;
+                                if let Some(prefix) = self.segment_override_for_op(1) {
+                                    write!(out, "{}:", prefix);
+                                }
                                 x.colorize(colors, out)
                             }
                         }
@@ -602,6 +631,9 @@ impl <T: std::fmt::Write> ShowContextual<u64, [Option<String>], T> for Instructi
                             },
                             x @ _ => {
                                 write!(out, ", ")?;
+                                if let Some(prefix) = self.segment_override_for_op(1) {
+                                    write!(out, "{}:", prefix);
+                                }
                                 x.colorize(colors, out)
                             }
                         }
diff --git a/src/lib.rs b/src/lib.rs
index 234735f..bc60dfd 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -172,6 +172,38 @@ pub enum Operand {
     Nothing
 }
 
+impl Operand {
+    pub fn is_memory(&self) -> bool {
+        match self {
+            Operand::DisplacementU32(_) |
+            Operand::DisplacementU64(_) |
+            Operand::RegDeref(_) |
+            Operand::RegDisp(_, _) |
+            Operand::RegScale(_, _) |
+            Operand::RegIndexBase(_, _) |
+            Operand::RegIndexBaseDisp(_, _, _) |
+            Operand::RegScaleDisp(_, _, _) |
+            Operand::RegIndexBaseScale(_, _, _) |
+            Operand::RegIndexBaseScaleDisp(_, _, _, _) => {
+                true
+            },
+            Operand::ImmediateI8(_) |
+            Operand::ImmediateU8(_) |
+            Operand::ImmediateI16(_) |
+            Operand::ImmediateU16(_) |
+            Operand::ImmediateU32(_) |
+            Operand::ImmediateI32(_) |
+            Operand::ImmediateU64(_) |
+            Operand::ImmediateI64(_) |
+            Operand::Register(_) |
+            Operand::Many(_) |
+            Operand::Nothing => {
+                false
+            }
+        }
+    }
+}
+
 #[cfg(feature="use-serde")]
 #[derive(Copy, Clone, Debug, Eq, PartialEq, Hash, Serialize, Deserialize)]
 pub enum RegisterBank {
@@ -189,7 +221,7 @@ pub enum RegisterBank {
     ST, MM,     // ST, MM regs (x87, mmx)
 }
 
-#[derive(Debug)]
+#[derive(Copy, Clone, Debug, Eq, PartialEq, Hash)]
 pub enum Segment {
     CS, DS, ES, FS, GS, SS
 }
@@ -421,6 +453,51 @@ impl Instruction {
             _ => false
         }
     }
+    pub fn segment_override_for_op(&self, op: u8) -> Option<Segment> {
+        match self.opcode {
+            Opcode::STOS => {
+                if op == 0 {
+                    Some(Segment::ES)
+                } else {
+                    None
+                }
+            }
+            Opcode::LODS => {
+                if op == 1 {
+                    Some(self.prefixes.segment)
+                } else {
+                    None
+                }
+            }
+            Opcode::MOVS => {
+                if op == 0 {
+                    Some(Segment::ES)
+                } else if op == 1 {
+                    Some(self.prefixes.segment)
+                } else {
+                    None
+                }
+            }
+            Opcode::CMPS => {
+                if op == 0 {
+                    Some(self.prefixes.segment)
+                } else if op == 1 {
+                    Some(Segment::ES)
+                } else {
+                    None
+                }
+            },
+            _ => {
+                // most operands are pretty simple:
+                if self.operands[op as usize].is_memory() &&
+                    self.prefixes.segment != Segment::DS {
+                    Some(self.prefixes.segment)
+                } else {
+                    None
+                }
+            }
+        }
+    }
 }
 
 #[derive(Debug)]
diff --git a/test/test.rs b/test/test.rs
index 1e73a2e..f8dadbb 100644
--- a/test/test.rs
+++ b/test/test.rs
@@ -83,6 +83,7 @@ fn test_mov() {
     test_display(&[0x0f, 0xb6, 0x06], "movzx eax, byte [rsi]");
     test_display(&[0x0f, 0xb7, 0x06], "movzx eax, word [rsi]");
     test_display(&[0x89, 0x55, 0x94], "mov [rbp - 0x6c], edx");
+    test_display(&[0x65, 0x4c, 0x89, 0x04, 0x25, 0xa8, 0x01, 0x00, 0x00], "mov gs:[0x1a8], r8");
 }
 
 #[test]